Source code
Revision control
Copy as Markdown
Other Tools
Test Info:
/* Any copyright is dedicated to the Public Domain.
const { SmartAssistEngine } = ChromeUtils.importESModule(
"moz-src:///browser/components/genai/SmartAssistEngine.sys.mjs"
);
const { sinon } = ChromeUtils.importESModule(
);
// Prefs
const PREF_API_KEY = "browser.ml.smartAssist.apiKey";
const PREF_ENDPOINT = "browser.ml.smartAssist.endpoint";
const PREF_MODEL = "browser.ml.smartAssist.model";
// Clean prefs after all tests
registerCleanupFunction(() => {
for (let pref of [PREF_API_KEY, PREF_ENDPOINT, PREF_MODEL]) {
if (Services.prefs.prefHasUserValue(pref)) {
Services.prefs.clearUserPref(pref);
}
}
});
add_task(async function test_createOpenAIEngine_uses_prefs_and_static_fields() {
Services.prefs.setStringPref(PREF_API_KEY, "test-key-123");
Services.prefs.setStringPref(PREF_MODEL, "gpt-fake");
const sb = sinon.createSandbox();
try {
const fakeEngine = {
runWithGenerator() {
throw new Error("not used");
},
};
const stub = sb
.stub(SmartAssistEngine, "_createEngine")
.resolves(fakeEngine);
const engine = await SmartAssistEngine.createOpenAIEngine();
Assert.strictEqual(
engine,
fakeEngine,
"Should return engine from _createEngine"
);
Assert.ok(stub.calledOnce, "_createEngine should be called once");
const opts = stub.firstCall.args[0];
Assert.equal(opts.apiKey, "test-key-123", "apiKey should come from pref");
Assert.equal(
opts.baseURL,
"baseURL should come from pref"
);
Assert.equal(opts.modelId, "gpt-fake", "modelId should come from pref");
} finally {
sb.restore();
}
});
add_task(async function test_fetchWithHistory_streams_and_forwards_args() {
const sb = sinon.createSandbox();
try {
let capturedArgs = null;
let capturedStreamOption = null;
// Fake async generator that yields three text chunks and one empty (ignored)
const fakeEngine = {
runWithGenerator({ streamOptions, args }) {
capturedArgs = args;
capturedStreamOption = streamOptions;
async function* gen() {
yield { text: "Hello" };
yield { text: " from" };
yield { text: " fake engine!" };
yield {}; // ignored by SmartAssistEngine
}
return gen();
},
};
sb.stub(SmartAssistEngine, "_createEngine").resolves(fakeEngine);
const messages = [
{ role: "system", content: "You are helpful" },
{ role: "user", content: "Hi there" },
];
// Collect streamed output
let acc = "";
for await (const t of SmartAssistEngine.fetchWithHistory(messages)) {
acc += t;
}
Assert.equal(
acc,
"Hello from fake engine!",
"Should concatenate streamed chunks"
);
Assert.deepEqual(
capturedArgs,
messages,
"Should forward messages as args to runWithGenerator()"
);
Assert.deepEqual(
capturedStreamOption.enabled,
true,
"Should enable streaming in runWithGenerator()"
);
} finally {
sb.restore();
}
});
add_task(
async function test_fetchWithHistory_propagates_engine_creation_rejection() {
const sb = sinon.createSandbox();
try {
const err = new Error("creation failed (generic)");
const stub = sb.stub(SmartAssistEngine, "_createEngine").rejects(err);
const messages = [{ role: "user", content: "Hi" }];
// Must CONSUME the async generator to trigger the rejection
const consume = async () => {
for await (const _message of SmartAssistEngine.fetchWithHistory(
messages
)) {
void _message;
}
};
await Assert.rejects(
consume(),
e => e === err,
"Should propagate the same error thrown by _createEngine"
);
Assert.ok(stub.calledOnce, "_createEngine should be called once");
} finally {
sb.restore();
}
}
);
add_task(async function test_fetchWithHistory_propagates_stream_error() {
const sb = sinon.createSandbox();
try {
const fakeEngine = {
runWithGenerator() {
async function* gen() {
yield { text: "partial" };
throw new Error("engine stream boom");
}
return gen();
},
};
sb.stub(SmartAssistEngine, "_createEngine").resolves(fakeEngine);
const consume = async () => {
let acc = "";
for await (const t of SmartAssistEngine.fetchWithHistory([
{ role: "user", content: "x" },
])) {
acc += t;
}
return acc;
};
await Assert.rejects(
consume(),
e => /engine stream boom/.test(e.message),
"Should propagate errors thrown during streaming"
);
} finally {
sb.restore();
}
});