Source code
Revision control
Copy as Markdown
Other Tools
Test Info:
/* Any copyright is dedicated to the Public Domain.
"use strict";
const { MLTestUtils } = ChromeUtils.importESModule(
);
/**
* Creates options for a model-free echo engine.
*
* @param {string} featureId - The feature requesting the engine.
* @param {string} [backend="onnx"] - The requested inference backend.
* @param {number} [numThreads=2] - The requested number of inference threads.
* @returns {{
* featureId: string,
* taskName: string,
* backend: string,
* numThreads: number,
* }} The engine options.
*/
function createOptions(featureId, backend = "onnx", numThreads = 2) {
return {
featureId,
taskName: "moz-echo",
backend,
numThreads,
};
}
/**
* Runs echo inference through real engines and resolves their mocked runtimes.
*
* @param {Array<MLEngine<any>>} engines - The engines to invoke.
* @param {Record<string, AttachmentMock>} remoteClients - The mocked Remote Settings clients.
* @returns {Promise<void>}
*/
async function runEchoEngines(engines, remoteClients) {
const messages = engines.map((_, index) => `Engine ${index}`);
const runs = engines.map((engine, index) =>
engine.run({ data: messages[index] })
);
await remoteClients["ml-onnx-runtime"].resolvePendingDownloads(
engines.length
);
const results = await Promise.all(runs);
Assert.deepEqual(
results.map(result => result.output.echo),
messages,
"Every real engine completes inference"
);
}
/**
* Restores engine interception and the ML test environment.
*
* @param {() => Promise<void>} environmentCleanup - Restores the ML test
* environment.
* @returns {Promise<void>}
*/
async function cleanupTest(environmentCleanup) {
try {
MLTestUtils.cleanupEngineCreationInterceptions();
} finally {
await EngineProcess.destroyMLEngine();
await environmentCleanup();
}
}
/**
* This test case ensures that multiple engine interceptions can be armed while
* unrelated engine creation continues normally.
*/
add_task(async function test_intercepts_multiple_engine_creations() {
const { cleanup, remoteClients } = await setup();
const originalGetEngine = MLEngineParent.prototype.getEngine;
const encoderCreation = MLTestUtils.interceptEngineCreation(
"formfill-encoder",
{
expectedOptions: { backend: "onnx", numThreads: 2 },
overrides: { numThreads: 1 },
}
);
const installedGetEngine = MLEngineParent.prototype.getEngine;
Assert.notEqual(
installedGetEngine,
originalGetEngine,
"Arming the first engine replaces getEngine"
);
const headCreation = MLTestUtils.interceptEngineCreation("formfill-head", {
expectedOptions: { backend: "onnx", numThreads: 2 },
overrides: { numThreads: 1 },
});
Assert.equal(
MLEngineParent.prototype.getEngine,
installedGetEngine,
"Arming another engine keeps the installed interceptor"
);
try {
const unrelatedEngine = await createEngine(
createOptions("about-inference")
);
const [encoderEngine, headEngine] = await Promise.all([
createEngine(createOptions("formfill-encoder")),
createEngine(createOptions("formfill-head")),
]);
Assert.equal(
MLEngineParent.prototype.getEngine,
originalGetEngine,
"Consuming the final armed engine restores getEngine"
);
Assert.equal(
(await encoderCreation).engine,
encoderEngine,
"The encoder interceptor resolves with the real engine"
);
Assert.equal(
(await headCreation).engine,
headEngine,
"The head interceptor resolves with the real engine"
);
Assert.equal(
unrelatedEngine.pipelineOptions.numThreads,
2,
"Unrelated engine options pass through unchanged"
);
Assert.equal(
encoderEngine.pipelineOptions.numThreads,
1,
"The encoder receives the thread override"
);
Assert.equal(
headEngine.pipelineOptions.numThreads,
1,
"The head receives the thread override"
);
const repeatedCreation = MLTestUtils.interceptEngineCreation(
"formfill-encoder",
{
expectedOptions: { backend: "onnx", numThreads: 2 },
overrides: { numThreads: 1 },
}
);
Assert.notEqual(
MLEngineParent.prototype.getEngine,
originalGetEngine,
"Rearming an engine reinstalls interception"
);
const repeatedEngine = await createEngine(
createOptions("formfill-encoder")
);
Assert.equal(
MLEngineParent.prototype.getEngine,
originalGetEngine,
"Consuming the rearmed engine restores getEngine"
);
Assert.equal(
(await repeatedCreation).engine,
repeatedEngine,
"The interceptor can be rearmed for the encoder"
);
Assert.equal(
repeatedEngine,
encoderEngine,
"Repeated creation reuses the real encoder engine"
);
await runEchoEngines(
[unrelatedEngine, encoderEngine, headEngine],
remoteClients
);
} finally {
await cleanupTest(cleanup);
}
});
/**
* This test case ensures that an engine cannot be armed twice before its
* expected creation occurs.
*/
add_task(async function test_rejects_duplicate_engine_interceptors() {
const { cleanup, remoteClients } = await setup();
const engineCreation = MLTestUtils.interceptEngineCreation(
"formfill-classification"
);
try {
Assert.throws(
() => MLTestUtils.interceptEngineCreation("formfill-classification"),
/already armed/,
"The same feature cannot be armed twice"
);
const engine = await createEngine(createOptions("formfill-classification"));
Assert.equal(
(await engineCreation).engine,
engine,
"The original interceptor remains armed"
);
await runEchoEngines([engine], remoteClients);
} finally {
await cleanupTest(cleanup);
}
});
/**
* This test case ensures that a WASM engine request can be replaced with native
* ONNX while still completing real inference.
*/
add_task(async function test_interceptor_overrides_wasm_with_native() {
const { cleanup } = await setup();
try {
if (!(await EngineProcess.requestIsNativeOnnxRuntimeAvailable())) {
info("Skipping because the native ONNX runtime is unavailable");
return;
}
const engineCreation = MLTestUtils.interceptEngineCreation(
"formfill-head",
{
expectedOptions: { backend: "onnx" },
overrides: { backend: "onnx-native" },
}
);
const engine = await createEngine(createOptions("formfill-head"));
const result = await engine.run({ data: "This gets echoed." });
Assert.equal(
(await engineCreation).engine,
engine,
"The interceptor resolves with the native engine"
);
Assert.equal(
engine.pipelineOptions.backend,
"onnx-native",
"The interceptor created a native ONNX engine"
);
Assert.equal(
result.output.echo,
"This gets echoed.",
"Inference completes through the native engine"
);
} finally {
await cleanupTest(cleanup);
}
});
/**
* This test case ensures that unexpected engine options fail the test task
* before engine creation proceeds.
*/
add_task(async function test_interceptor_validates_options() {
const { cleanup } = await setup();
const engineCreation = MLTestUtils.interceptEngineCreation(
"formfill-encoder",
{
expectedOptions: { backend: "onnx-native", numThreads: 2 },
}
);
const interceptedRejection = engineCreation.catch(error => error);
try {
await Assert.rejects(
createEngine(createOptions("formfill-encoder")),
/Expected engine option "backend" for "formfill-encoder"/,
"Engine creation rejects the unexpected backend"
);
Assert.stringContains(
(await interceptedRejection).message,
'Expected engine option "backend" for "formfill-encoder"',
"The interceptor rejects the same unexpected backend"
);
} finally {
await cleanupTest(cleanup);
}
});
/**
* This test case ensures cleanup restores engine creation before reporting an
* unconsumed interceptor.
*/
add_task(function test_cleanup_rejects_armed_interceptors_and_restores() {
const originalGetEngine = MLEngineParent.prototype.getEngine;
MLTestUtils.interceptEngineCreation("formfill-classification");
Assert.notEqual(
MLEngineParent.prototype.getEngine,
originalGetEngine,
"Arming an engine replaces getEngine"
);
Assert.throws(
MLTestUtils.cleanupEngineCreationInterceptions,
/still armed for: formfill-classification/,
"Cleanup rejects an unconsumed engine interceptor"
);
Assert.equal(
MLEngineParent.prototype.getEngine,
originalGetEngine,
"Cleanup restores getEngine before rejecting"
);
});