1 files changed, 155 insertions, 0 deletions
diff --git a/toolkit/components/ml/tests/browser/head.js b/toolkit/components/ml/tests/browser/head.js
new file mode 100644
index 0000000000..99d27ce18a
--- /dev/null
+++ b/toolkit/components/ml/tests/browser/head.js
@@ -0,0 +1,155 @@
+/* Any copyright is dedicated to the Public Domain.
+   http://creativecommons.org/publicdomain/zero/1.0/ */
+
+/// <reference path="../../../../../toolkit/components/translations/tests/browser/shared-head.js" />
+
+"use strict";
+
+/**
+ * @type {import("../../content/SummarizerModel.sys.mjs")}
+ */
+const { SummarizerModel } = ChromeUtils.importESModule(
+  "chrome://global/content/ml/SummarizerModel.sys.mjs"
+);
+
+/**
+ * @type {import("../../actors/MLEngineParent.sys.mjs")}
+ */
+const { MLEngineParent } = ChromeUtils.importESModule(
+  "resource://gre/actors/MLEngineParent.sys.mjs"
+);
+
+// This test suite shares some utility functions with translations as they work in a very
+// similar fashion. Eventually, the plan is to unify these two components.
+Services.scriptloader.loadSubScript(
+  "chrome://mochitests/content/browser/toolkit/components/translations/tests/browser/shared-head.js",
+  this
+);
+
+function getDefaultModelRecords() {
+  return [
+    {
+      name: "summarizer-model",
+      version: SummarizerModel.MODEL_MAJOR_VERSION + ".0",
+    },
+  ];
+}
+
+function getDefaultWasmRecords() {
+  return [
+    {
+      name: "inference-engine",
+      version: MLEngineParent.WASM_MAJOR_VERSION + ".0",
+    },
+  ];
+}
+
+/**
+ * Creates a local RemoteSettingsClient for use within tests.
+ *
+ * @param {boolean} autoDownloadFromRemoteSettings
+ * @param {object[]} models
+ * @returns {AttachmentMock}
+ */
+async function createMLModelsRemoteClient(
+  autoDownloadFromRemoteSettings,
+  models = getDefaultModelRecords()
+) {
+  const { RemoteSettings } = ChromeUtils.importESModule(
+    "resource://services-settings/remote-settings.sys.mjs"
+  );
+  const mockedCollectionName = "test-ml-models";
+  const client = RemoteSettings(
+    `${mockedCollectionName}-${_remoteSettingsMockId++}`
+  );
+  const metadata = {};
+  await client.db.clear();
+  await client.db.importChanges(
+    metadata,
+    Date.now(),
+    models.map(({ name, version }) => ({
+      id: crypto.randomUUID(),
+      name,
+      version,
+      last_modified: Date.now(),
+      schema: Date.now(),
+      attachment: {
+        hash: `${crypto.randomUUID()}`,
+        size: `123`,
+        filename: name,
+        location: `main-workspace/ml-models/${crypto.randomUUID()}.bin`,
+        mimetype: "application/octet-stream",
+      },
+    }))
+  );
+
+  return createAttachmentMock(
+    client,
+    mockedCollectionName,
+    autoDownloadFromRemoteSettings
+  );
+}
+
+async function createAndMockMLRemoteSettings({
+  models = getDefaultModelRecords(),
+  autoDownloadFromRemoteSettings = false,
+} = {}) {
+  const remoteClients = {
+    models: await createMLModelsRemoteClient(
+      autoDownloadFromRemoteSettings,
+      models
+    ),
+    wasm: await createMLWasmRemoteClient(autoDownloadFromRemoteSettings),
+  };
+
+  MLEngineParent.mockRemoteSettings(remoteClients.wasm.client);
+  SummarizerModel.mockRemoteSettings(remoteClients.models.client);
+
+  return {
+    async removeMocks() {
+      await remoteClients.models.client.attachments.deleteAll();
+      await remoteClients.models.client.db.clear();
+      await remoteClients.wasm.client.attachments.deleteAll();
+      await remoteClients.wasm.client.db.clear();
+
+      MLEngineParent.removeMocks();
+      SummarizerModel.removeMocks();
+    },
+    remoteClients,
+  };
+}
+
+/**
+ * Creates a local RemoteSettingsClient for use within tests.
+ *
+ * @param {boolean} autoDownloadFromRemoteSettings
+ * @returns {AttachmentMock}
+ */
+async function createMLWasmRemoteClient(autoDownloadFromRemoteSettings) {
+  const { RemoteSettings } = ChromeUtils.importESModule(
+    "resource://services-settings/remote-settings.sys.mjs"
+  );
+  const mockedCollectionName = "test-translation-wasm";
+  const client = RemoteSettings(
+    `${mockedCollectionName}-${_remoteSettingsMockId++}`
+  );
+  const metadata = {};
+  await client.db.clear();
+  await client.db.importChanges(
+    metadata,
+    Date.now(),
+    getDefaultWasmRecords().map(({ name, version }) => ({
+      id: crypto.randomUUID(),
+      name,
+      version,
+      last_modified: Date.now(),
+      schema: Date.now(),
+    }))
+  );
+
+  return createAttachmentMock(
+    client,
+    mockedCollectionName,
+    autoDownloadFromRemoteSettings
+  );
+}