Source code

Revision control

Copy as Markdown

Other Tools

Test Info: Warnings

/* This Source Code Form is subject to the terms of the Mozilla Public
* License, v. 2.0. If a copy of the MPL was not distributed with this
* file, You can obtain one at http://mozilla.org/MPL/2.0/. */
"use strict";
const lazy = {};
ChromeUtils.defineESModuleGetters(lazy, {
AIWindow:
"moz-src:///browser/components/aiwindow/ui/modules/AIWindow.sys.mjs",
IntentClassifier:
"moz-src:///browser/components/aiwindow/models/IntentClassifier.sys.mjs",
UrlbarProviderAiChat:
"moz-src:///browser/components/urlbar/UrlbarProviderAiChat.sys.mjs",
UrlbarSearchUtils:
"moz-src:///browser/components/urlbar/UrlbarSearchUtils.sys.mjs",
});
const fakeEngine = {
run({ args: [[query]] }) {
const searchKeywords = ["search", "find"];
const formattedPrompt = query.toLowerCase();
const isSearch = searchKeywords.some(keyword =>
formattedPrompt.includes(keyword)
);
// Simulate model confidence scores
if (isSearch) {
return [
{ label: "search", score: 0.95 },
{ label: "chat", score: 0.05 },
];
}
return [
{ label: "chat", score: 0.95 },
{ label: "search", score: 0.05 },
];
},
};
let gEngineStub;
add_task(async function setup() {
gEngineStub = sinon
.stub(lazy.IntentClassifier, "_createEngine")
.resolves(fakeEngine);
registerCleanupFunction(() => {
sinon.restore();
});
await lazy.UrlbarSearchUtils.init();
// Set required prefs
Services.prefs.setBoolPref("browser.ml.enable", true);
});
add_task(async function test_isActive() {
let aiWindowStub = sinon
.stub(lazy.AIWindow, "isAIWindowActiveAndEnabled")
.returns(false);
const provider = new lazy.UrlbarProviderAiChat();
let controller = UrlbarTestUtils.newMockController();
Assert.ok(
!(await provider.isActive(createContext("hello"), controller)),
"Provider is not active as this is not a Smart Window"
);
aiWindowStub.returns(true);
Assert.ok(
await provider.isActive(createContext("hello"), controller),
"Provider is active"
);
Assert.ok(
!(await provider.isActive(
createContext("hello world", {
searchMode: { engineName: "testEngine" },
}),
controller
)),
"Provider is not active as we're in Search Mode"
);
Assert.ok(
!(await provider.isActive(createContext("hi"), controller)),
"Provider is not active as prompt is too short"
);
});
add_task(async function test_startQuery() {
const provider = new lazy.UrlbarProviderAiChat();
let engine = lazy.UrlbarSearchUtils.getDefaultEngine();
let icon = await engine.getIconURL();
const TESTS = [
{
desc: "Urlbar for 'chat' intent returns non-heuristic with suggestedIndex",
sapName: "urlbar",
input: "tell me a joke",
expectedResults: [
{
heuristic: false,
type: UrlbarUtils.RESULT_TYPE.AI_CHAT,
suggestedIndex: 1,
query: "tell me a joke",
icon: lazy.UrlbarProviderAiChat.CHAT_ICON_URL,
title: "tell me a joke",
},
],
},
{
desc: "Urlbar for 'search' intent doesn't return any result",
sapName: "urlbar",
input: "search for cute cat pictures",
expectedResults: [],
},
{
desc: "Urlbar for 'navigate' intent doesn't return any result",
sapName: "urlbar",
input: "mozilla.org",
expectedResults: [],
},
{
desc: "Smartbar for 'search' intent returns a non heuristic result",
sapName: "smartbar",
input: "search for cute cat pictures",
expectedResults: [
{
heuristic: false,
type: UrlbarUtils.RESULT_TYPE.AI_CHAT,
suggestedIndex: 1,
query: "search for cute cat pictures",
icon: lazy.UrlbarProviderAiChat.CHAT_ICON_URL,
title: "search for cute cat pictures",
},
],
},
{
desc: "Smartbar for 'chat' intent returns heuristic and search result",
sapName: "smartbar",
input: "tell me a joke",
expectedResults: [
{
heuristic: true,
type: UrlbarUtils.RESULT_TYPE.AI_CHAT,
suggestedIndex: undefined,
query: "tell me a joke",
icon: lazy.UrlbarProviderAiChat.CHAT_ICON_URL,
title: "tell me a joke",
},
{
heuristic: false,
type: UrlbarUtils.RESULT_TYPE.SEARCH,
suggestedIndex: undefined,
query: "tell me a joke",
icon,
engine: engine.name,
title: "tell me a joke",
},
],
},
{
desc: "Smartbar for 'navigate' intent doesn't return any result",
sapName: "smartbar",
input: "mozilla.org",
expectedResults: [],
},
];
for (const { desc, sapName, input, expectedResults } of TESTS) {
info(desc);
let added = [];
await provider.startQuery(
createContext(input, {
sapName,
}),
(_provider, result) => {
added.push(result);
}
);
Assert.equal(
added.length,
expectedResults.length,
"Check number of results"
);
for (let i = 0; i < expectedResults.length; i++) {
let result = added[i];
let expected = expectedResults[i];
for (let prop in expected) {
let checkFn =
expected[prop] === undefined
? Assert.strictEqual.bind(Assert)
: Assert.equal.bind(Assert);
let checkObj = ["engine", "icon", "query", "title"].includes(prop)
? result.payload
: result;
checkFn(
checkObj[prop],
expected[prop],
`Check result ${i} property ${prop}`
);
}
}
}
});
add_task(async function test_intentClassifierThrows() {
const provider = new lazy.UrlbarProviderAiChat();
info("Stub the engine to throw, we should fallback to 'search' intent.");
const throwingEngine = {
run() {
throw new Error("Intent classifier failure");
},
};
gEngineStub.resolves(throwingEngine);
let added = [];
await provider.startQuery(
createContext("tell me a joke", {
sapName: "smartbar",
}),
(_provider, result) => {
added.push(result);
}
);
Assert.equal(
added.length,
1,
"One result is returned, because we fallback to 'search' intent"
);
let result = added[0];
Assert.equal(
result.payload.query,
"tell me a joke",
"The result has the expected query"
);
Assert.equal(
result.type,
UrlbarUtils.RESULT_TYPE.AI_CHAT,
"The result is a chat result"
);
Assert.equal(
result.suggestedIndex,
1,
"The result is non-heuristic (suggestedIndex is set)"
);
Assert.equal(
result.heuristic,
false,
"The result is non-heuristic (heuristic is false)"
);
});