UNPKG

worker-testbed

Version:

The AI-ready testbed which use Worker Threads to create isolated vm context

85 lines (82 loc) 2.6 kB
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 };