434 lines
13 KiB
JavaScript
434 lines
13 KiB
JavaScript
/* Any copyright is dedicated to the Public Domain.
|
|
* http://creativecommons.org/publicdomain/zero/1.0/ */
|
|
|
|
"use strict";
|
|
const { ModelHubProvider } = ChromeUtils.importESModule(
|
|
"resource://gre/modules/addons/ModelHubProvider.sys.mjs"
|
|
);
|
|
ChromeUtils.defineESModuleGetters(this, {
|
|
addonIdToEngineId: "chrome://global/content/ml/Utils.sys.mjs",
|
|
engineIdToAddonId: "chrome://global/content/ml/Utils.sys.mjs",
|
|
isAddonEngineId: "chrome://global/content/ml/Utils.sys.mjs",
|
|
sinon: "resource://testing-common/Sinon.sys.mjs",
|
|
});
|
|
createAppInfo("xpcshell@tests.mozilla.org", "XPCShell", "1");
|
|
AddonTestUtils.init(this);
|
|
|
|
const LOCAL_MODEL_MANAGEMENT_ENABLED_PREF =
|
|
"extensions.htmlaboutaddons.local_model_management";
|
|
|
|
function ensureBrowserDelayedStartupFinished() {
|
|
// ModelHubProvider does not register itself until the application startup
|
|
// has been completed, and so we simulate that by firing this notification.
|
|
Services.obs.notifyObservers(null, "browser-delayed-startup-finished");
|
|
}
|
|
|
|
add_setup(async () => {
|
|
await promiseStartupManager();
|
|
});
|
|
|
|
add_task(
|
|
{
|
|
pref_set: [[LOCAL_MODEL_MANAGEMENT_ENABLED_PREF, false]],
|
|
},
|
|
async function test_modelhub_provider_disabled() {
|
|
ensureBrowserDelayedStartupFinished();
|
|
ok(
|
|
!AddonManager.hasProvider("ModelHubProvider"),
|
|
"Expect no ModelHubProvider to be registered"
|
|
);
|
|
}
|
|
);
|
|
|
|
add_task(
|
|
{
|
|
pref_set: [[LOCAL_MODEL_MANAGEMENT_ENABLED_PREF, true]],
|
|
},
|
|
async function test_modelhub_provider_enabled() {
|
|
ensureBrowserDelayedStartupFinished();
|
|
ok(
|
|
AddonManager.hasProvider("ModelHubProvider"),
|
|
"Expect ModelHubProvider to be registered"
|
|
);
|
|
}
|
|
);
|
|
|
|
add_task(
|
|
{
|
|
pref_set: [[LOCAL_MODEL_MANAGEMENT_ENABLED_PREF, true]],
|
|
},
|
|
async function test_modelhub_provider_addon_wrappers() {
|
|
let sandbox = sinon.createSandbox();
|
|
ModelHubProvider.clearAddonCache();
|
|
// Sanity checks.
|
|
ok(
|
|
AddonManager.hasProvider("ModelHubProvider"),
|
|
"Expect ModelHubProvider to be registered"
|
|
);
|
|
Assert.ok(
|
|
ModelHubProvider.modelHub,
|
|
"Expect modelHub instance to be found"
|
|
);
|
|
|
|
const fakeAddonIds = [
|
|
"addon1-using-model@test-extension",
|
|
"addon2-using-model@test-extension",
|
|
];
|
|
|
|
const mockModels = [
|
|
{
|
|
name: "model-hub.mozilla.org/org1/model-mock-1",
|
|
revision: "mockRevision1",
|
|
engineIds: [],
|
|
},
|
|
{
|
|
name: "huggingface.co/org2/model-mock-2",
|
|
revision: "mockRevision2",
|
|
engineIds: [],
|
|
},
|
|
];
|
|
const mockListFilesResult = {
|
|
metadata: {
|
|
totalSize: 2048,
|
|
lastUsed: new Date("2023-10-01T12:00:00Z"),
|
|
updateDate: 0,
|
|
// This is retuned by the first call to the listFiles
|
|
// stub and then used to determine what are the expected
|
|
// usedByFirefoxFeatures and usedByAddonIds properties
|
|
// on the wrapper.
|
|
engineIds: [
|
|
"about-inference",
|
|
"non-existing-feature",
|
|
...fakeAddonIds.map(addonId => addonIdToEngineId(addonId)),
|
|
],
|
|
},
|
|
};
|
|
|
|
const mockModelShortNames = ["model-mock-1", "model-mock-2"];
|
|
const mockModelsHomepageURLs = [
|
|
"https://huggingface.co/org1/model-mock-1/",
|
|
"https://huggingface.co/org2/model-mock-2/",
|
|
];
|
|
|
|
const listModelsStub = sandbox
|
|
.stub(ModelHubProvider.modelHub, "listModels")
|
|
.resolves(mockModels);
|
|
|
|
const listFilesStub = sandbox
|
|
.stub(ModelHubProvider.modelHub, "listFiles")
|
|
.onFirstCall()
|
|
.resolves(mockListFilesResult)
|
|
.onSecondCall()
|
|
.resolves({
|
|
metadata: {
|
|
...mockListFilesResult.metadata,
|
|
// Setting engineIds to undefined to confirm that it is
|
|
// going to be set it to an empty array.
|
|
engineIds: undefined,
|
|
},
|
|
});
|
|
|
|
const getOwnerIcon = sandbox
|
|
.stub(ModelHubProvider.modelHub, "getOwnerIcon")
|
|
.resolves("chrome://mozapps/skin/extensions/extensionGeneric.svg");
|
|
|
|
const deleteModels = sandbox
|
|
.stub(ModelHubProvider.modelHub, "deleteModels")
|
|
.resolves();
|
|
|
|
const modelWrappers = await AddonManager.getAddonsByTypes(["mlmodel"]);
|
|
|
|
// Check that the stubs were called the expected number of times.
|
|
Assert.equal(
|
|
listModelsStub.callCount,
|
|
1,
|
|
"listModels() getting all models only once"
|
|
);
|
|
Assert.equal(
|
|
listFilesStub.callCount,
|
|
mockModels.length,
|
|
"listFiles() getting files and file metadata once for each model"
|
|
);
|
|
|
|
Assert.equal(
|
|
getOwnerIcon.callCount,
|
|
mockModels.length,
|
|
"getOwnerIcon() getting image blob once for each model"
|
|
);
|
|
|
|
// Verify that the listFiles was called with the expected arguments.
|
|
for (let i = 0; i < mockModels.length; i++) {
|
|
// ListFiles only has one argument, which is an config object
|
|
const callArgs = listFilesStub.getCall(i).args[0];
|
|
|
|
Assert.ok(callArgs, `listFiles call ${i} received arguments`);
|
|
// Compare the model name and revision arguments to the mockModels.
|
|
Assert.equal(
|
|
callArgs.model,
|
|
mockModels[i].name,
|
|
`Correct model name for call ${i}`
|
|
);
|
|
Assert.equal(
|
|
callArgs.revision,
|
|
mockModels[i].revision,
|
|
`Correct revision for call ${i}`
|
|
);
|
|
}
|
|
|
|
Assert.equal(
|
|
modelWrappers.length,
|
|
mockModels.length,
|
|
"Got the expected number of model AddonWrapper instances"
|
|
);
|
|
|
|
for (const [idx, modelWrapper] of modelWrappers.entries()) {
|
|
const { name, revision } = mockModels[idx];
|
|
const { engineIds } = mockListFilesResult.metadata;
|
|
// The first call to the listFiles stub is expected to include
|
|
// engineIds, whereare the second one is expected to not be
|
|
// including it.
|
|
const usedByFirefoxFeatures =
|
|
idx === 0
|
|
? engineIds.filter(engineId => !isAddonEngineId(engineId))
|
|
: [];
|
|
const usedByAddonIds =
|
|
idx === 0
|
|
? engineIds
|
|
.filter(engineId => isAddonEngineId(engineId))
|
|
.map(engineId => engineIdToAddonId(engineId))
|
|
: [];
|
|
verifyModelAddonWrapper(modelWrapper, {
|
|
model: name,
|
|
name: mockModelShortNames[idx],
|
|
version: revision,
|
|
lastUsed: mockListFilesResult.metadata.lastUsed,
|
|
totalSize: mockListFilesResult.metadata.totalSize,
|
|
modelHomepageURL: mockModelsHomepageURLs[idx],
|
|
usedByFirefoxFeatures,
|
|
usedByAddonIds,
|
|
});
|
|
}
|
|
|
|
// Verify that the ModelHubProvider.getAddonsByTypes
|
|
// doesn't return any entry if mlmodel isn't explicitly
|
|
// requested.
|
|
Assert.deepEqual(
|
|
(await AddonManager.getAddonsByTypes()).filter(
|
|
addon => addon.type === "mlmodel"
|
|
),
|
|
[],
|
|
"Expect no mlmodel results with getAddonsByTypes()"
|
|
);
|
|
|
|
Assert.deepEqual(
|
|
(await AddonManager.getAddonsByTypes([])).filter(
|
|
addon => addon.type === "mlmodel"
|
|
),
|
|
[],
|
|
"Expect no mlmodel results with getAddonsByTypes([])"
|
|
);
|
|
|
|
Assert.deepEqual(
|
|
(await AddonManager.getAddonsByTypes(["extension"])).filter(
|
|
addon => addon.type === "mlmodel"
|
|
),
|
|
[],
|
|
"Expect no mlmodel result with getAddonsByTypes(['extension'])"
|
|
);
|
|
|
|
Assert.equal(
|
|
await AddonManager.getAddonByID(modelWrappers[0].id),
|
|
modelWrappers[0],
|
|
`Got the expected result from getAddonByID for ${modelWrappers[0].id}`
|
|
);
|
|
|
|
// Selecting first model wrapper to test uninstall.
|
|
const modelWrapper = modelWrappers[0];
|
|
const uninstallPromise = AddonTestUtils.promiseAddonEvent("onUninstalled");
|
|
await modelWrapper.uninstall();
|
|
|
|
const [uninstalled] = await uninstallPromise;
|
|
equal(
|
|
uninstalled,
|
|
modelWrapper,
|
|
"onUninstalled was called with that wrapper"
|
|
);
|
|
|
|
// We expect getAddonByID for the removed model to not be found.
|
|
Assert.equal(
|
|
await AddonManager.getAddonByID(modelWrappers[0].id),
|
|
null,
|
|
`Got no model wrapper from getAddonByID for uninstalled ${modelWrappers[0].id}`
|
|
);
|
|
|
|
// We expect getAddonByID for the non removed model to still be found.
|
|
Assert.equal(
|
|
await AddonManager.getAddonByID(modelWrappers[1].id),
|
|
modelWrappers[1],
|
|
`Got the expected result from getAddonByID for ${modelWrappers[1].id}`
|
|
);
|
|
|
|
Assert.equal(
|
|
deleteModels.callCount,
|
|
1,
|
|
"Got the expected number of ModelHub.deleteModels() method calls"
|
|
);
|
|
|
|
Assert.deepEqual(
|
|
deleteModels.firstCall.args,
|
|
[
|
|
{
|
|
model: mockModels[0].name,
|
|
revision: mockModels[0].revision,
|
|
deletedBy: "about:addons",
|
|
},
|
|
],
|
|
"Got the expected arguments in the ModelHub.deleteModels() method call"
|
|
);
|
|
|
|
// Reset all sinon stubs.
|
|
sandbox.restore();
|
|
|
|
function verifyModelAddonWrapper(modelWrapper, expected) {
|
|
const {
|
|
name,
|
|
model,
|
|
version,
|
|
lastUsed,
|
|
modelHomepageURL,
|
|
usedByFirefoxFeatures,
|
|
usedByAddonIds,
|
|
} = expected;
|
|
info(`Verify model addon wrapper for ${name}:${version}`);
|
|
const expectedId = ModelHubProvider.getWrapperIdForModel({
|
|
name: model,
|
|
revision: version,
|
|
});
|
|
Assert.equal(modelWrapper.id, expectedId, "Got the expected id");
|
|
Assert.equal(modelWrapper.type, "mlmodel", "Got the expected type");
|
|
Assert.equal(
|
|
modelWrapper.permissions,
|
|
AddonManager.PERM_CAN_UNINSTALL,
|
|
"Got the expected permissions"
|
|
);
|
|
Assert.equal(modelWrapper.model, model, "Got the expected name patch");
|
|
Assert.equal(modelWrapper.name, name, "Got the expected name");
|
|
Assert.equal(modelWrapper.version, version, "Got the expected version");
|
|
Assert.equal(
|
|
modelWrapper.lastUsed.toISOString(),
|
|
lastUsed.toISOString(),
|
|
"Got the expected lastUsed"
|
|
);
|
|
Assert.equal(
|
|
modelWrapper.totalSize,
|
|
expected.totalSize,
|
|
"Got the expected file size"
|
|
);
|
|
Assert.equal(
|
|
modelWrapper.isActive,
|
|
true,
|
|
"Expect model AddonWrapper to be active"
|
|
);
|
|
Assert.equal(
|
|
modelWrapper.isCompatible,
|
|
true,
|
|
"Expect model AddonWrapper to be compatible"
|
|
);
|
|
Assert.equal(
|
|
modelWrapper.modelHomepageURL,
|
|
modelHomepageURL,
|
|
"Got the expect model homepage URL"
|
|
);
|
|
Assert.deepEqual(
|
|
modelWrapper.usedByFirefoxFeatures,
|
|
usedByFirefoxFeatures,
|
|
"Got the expected engineIds listed in usedByFirefoxFeatures"
|
|
);
|
|
Assert.deepEqual(
|
|
modelWrapper.usedByAddonIds,
|
|
usedByAddonIds,
|
|
"Got the expected addon ids listed in usedByAddonIds"
|
|
);
|
|
}
|
|
}
|
|
);
|
|
|
|
add_task(
|
|
{
|
|
pref_set: [[LOCAL_MODEL_MANAGEMENT_ENABLED_PREF, true]],
|
|
},
|
|
async function test_modelhub_resets_cache_on_refresh() {
|
|
let sandbox = sinon.createSandbox();
|
|
ModelHubProvider.clearAddonCache();
|
|
|
|
const mockModels = [
|
|
{
|
|
name: "model-hub.mozilla.org/org1/model-mock-1",
|
|
revision: "mockRevision1",
|
|
engineIds: [],
|
|
},
|
|
{
|
|
name: "huggingface.co/org2/model-mock-2",
|
|
revision: "mockRevision2",
|
|
engineIds: [],
|
|
},
|
|
];
|
|
|
|
const mockListFilesResult = {
|
|
metadata: {
|
|
totalSize: 2048,
|
|
lastUsed: new Date("2023-10-01T12:00:00Z"),
|
|
updateDate: 0,
|
|
engineIds: ["about-inference", "non-existing-feature"],
|
|
},
|
|
};
|
|
|
|
sandbox
|
|
.stub(ModelHubProvider.modelHub, "listModels")
|
|
.onFirstCall()
|
|
.resolves(mockModels)
|
|
.onSecondCall()
|
|
.resolves([]);
|
|
|
|
sandbox
|
|
.stub(ModelHubProvider.modelHub, "listFiles")
|
|
.onFirstCall()
|
|
.resolves(mockListFilesResult)
|
|
.onSecondCall()
|
|
.resolves({
|
|
metadata: {
|
|
...mockListFilesResult.metadata,
|
|
// Setting engineIds to undefined to confirm that it is
|
|
// going to be set it to an empty array.
|
|
engineIds: undefined,
|
|
},
|
|
});
|
|
|
|
sandbox
|
|
.stub(ModelHubProvider.modelHub, "getOwnerIcon")
|
|
.resolves("chrome://mozapps/skin/extensions/extensionGeneric.svg");
|
|
|
|
sandbox.stub(ModelHubProvider.modelHub, "deleteModels").resolves();
|
|
|
|
// First call to the provider should populate the cache with the models
|
|
// returned by `listModels()`.
|
|
let modelWrappers = await AddonManager.getAddonsByTypes(["mlmodel"]);
|
|
Assert.equal(
|
|
modelWrappers.length,
|
|
mockModels.length,
|
|
"Got the expected number of model AddonWrapper instances"
|
|
);
|
|
|
|
// Second call should clear the cache before adding the models from the
|
|
// `ModelHub`. In this case, `listModels()` will return an empty array so
|
|
// we should expect no model wrapper.
|
|
modelWrappers = await AddonManager.getAddonsByTypes(["mlmodel"]);
|
|
Assert.equal(
|
|
modelWrappers.length,
|
|
0,
|
|
"Got the expected number of model AddonWrapper instances after refresh"
|
|
);
|
|
|
|
sandbox.restore();
|
|
}
|
|
);
|