worker-testbed/build/index.mjs

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

import { isMainThread, parentPort, workerData, Worker } from 'worker_threads';
import { fileURLToPath } from 'url';
import { BehaviorSubject, Subject, ToolRegistry, singleshot, getErrorMessage, createAwaiter, sleep } 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", async ({ status, msg }) => {
            if (status === "pass") {
                test.pass(msg);
            }
            else if (status === "fail") {
                test.fail(msg);
                await sleep(100);
            }
            isFinished = true;
            worker.terminate();
            resolve();
        });
        worker.on("error", async (err) => {
            test.fail(`Worker error: ${getErrorMessage(err)}`);
            await sleep(100);
            resolve();
        });
        worker.on("exit", async (code) => {
            if (isFinished) {
                return;
            }
            if (code !== 0) {
                test.fail(`Worker stopped with exit code ${code}`);
                await sleep(100);
                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");
    }
    const finishTest = singleshot((status, msg) => {
        parentPort.postMessage({ status, msg });
        process.exit(0);
    });
    process.on('uncaughtException', function (err) {
        // fail current test in worker
        finishTest("fail", `Uncaught exception: ${getErrorMessage(err)}`);
    });
    process.on('unhandledRejection', (error) => {
        throw error;
    });
    testRegistry.get(workerData.testName).cb({
        pass: (msg) => finishTest("pass", msg),
        fail: (msg) => finishTest("fail", msg),
    });
});

export { run, test };