worker-testbed
Version:
The AI-ready testbed which use Worker Threads to create isolated vm context
85 lines (82 loc) • 2.6 kB
JavaScript
import { isMainThread, parentPort, workerData, Worker } from 'worker_threads';
import { fileURLToPath } from 'url';
import { BehaviorSubject, Subject, ToolRegistry, singleshot, createAwaiter, getErrorMessage } from 'functools-kit';
import tape from 'tape';
const workerFileSubject = new BehaviorSubject();
const finishSubject = new Subject();
let testRegistry = new ToolRegistry("workertest");
let testCounter = 0;
const waitForFile = async () => {
if (workerFileSubject.data) {
return workerFileSubject.data;
}
return await workerFileSubject.toPromise();
};
class TestWrapper {
testName;
cb;
constructor(testName, cb) {
this.testName = testName;
this.cb = cb;
}
}
const test = async (testName, cb) => {
if (!isMainThread) {
testRegistry = testRegistry.register(testName, new TestWrapper(testName, cb));
return;
}
const workerFile = await waitForFile();
testCounter += 1;
tape(testName, async (test) => {
const [awaiter, { resolve }] = createAwaiter();
let isFinished = false;
const worker = new Worker(workerFile, {
workerData: { testName },
});
worker.once("message", ({ status, msg }) => {
if (status === "pass") {
test.pass(msg);
}
else if (status === "fail") {
test.fail(msg);
}
isFinished = true;
worker.terminate();
resolve();
});
worker.on("error", (err) => {
test.fail(`Worker error: ${getErrorMessage(err)}`);
resolve();
});
worker.on("exit", (code) => {
if (isFinished) {
return;
}
if (code !== 0) {
test.fail(`Worker stopped with exit code ${code}`);
resolve();
}
});
{
await awaiter;
testCounter -= 1;
await finishSubject.next();
}
});
};
const run = singleshot(async (__filename, cb = () => { }) => {
if (isMainThread) {
await workerFileSubject.next(fileURLToPath(__filename));
await finishSubject.filter(() => testCounter === 0).toPromise();
cb();
return;
}
if (!parentPort) {
throw new Error("workertest parentPort is null");
}
testRegistry.get(workerData.testName).cb({
pass: (msg) => parentPort.postMessage({ status: "pass", msg }),
fail: (msg) => parentPort.postMessage({ status: "fail", msg }),
});
});
export { run, test };