1 files changed, 48 insertions, 18 deletions

diff --git a/browser/components/search/test/unit/test_search_telemetry_categorization_sync.js b/browser/components/search/test/unit/test_search_telemetry_categorization_sync.js
index 423ee0a81d..40d38efbba 100644
--- a/browser/components/search/test/unit/test_search_telemetry_categorization_sync.js
+++ b/browser/components/search/test/unit/test_search_telemetry_categorization_sync.js
@@ -16,16 +16,34 @@ ChromeUtils.defineESModuleGetters(this, {
   TestUtils: "resource://testing-common/TestUtils.sys.mjs",
 });
 
+ChromeUtils.defineLazyGetter(this, "gCryptoHash", () => {
+  return Cc["@mozilla.org/security/hash;1"].createInstance(Ci.nsICryptoHash);
+});
+
+function convertDomainsToHashes(domainsToCategories) {
+  let newObj = {};
+  for (let [key, value] of Object.entries(domainsToCategories)) {
+    gCryptoHash.init(gCryptoHash.SHA256);
+    let bytes = new TextEncoder().encode(key);
+    gCryptoHash.update(bytes, key.length);
+    let hash = gCryptoHash.finish(true);
+    newObj[hash] = value;
+  }
+  return newObj;
+}
+
 async function waitForDomainToCategoriesUpdate() {
   return TestUtils.topicObserved("domain-to-categories-map-update-complete");
 }
 
-async function mockRecordWithCachedAttachment({ id, version, filename }) {
+async function mockRecordWithCachedAttachment({
+  id,
+  version,
+  filename,
+  mapping,
+}) {
   // Get the bytes of the file for the hash and size for attachment metadata.
-  let data = await IOUtils.readUTF8(
-    PathUtils.join(do_get_cwd().path, filename)
-  );
-  let buffer = new TextEncoder().encode(data).buffer;
+  let buffer = new TextEncoder().encode(JSON.stringify(mapping)).buffer;
   let stream = Cc["@mozilla.org/io/arraybuffer-input-stream;1"].createInstance(
     Ci.nsIArrayBufferInputStream
   );
@@ -73,21 +91,33 @@ const RECORDS = {
     id: RECORD_A_ID,
     version: 1,
     filename: "domain_category_mappings_1a.json",
+    mapping: convertDomainsToHashes({
+      "example.com": [1, 100],
+    }),
   },
   record1b: {
     id: RECORD_B_ID,
     version: 1,
     filename: "domain_category_mappings_1b.json",
+    mapping: convertDomainsToHashes({
+      "example.org": [2, 90],
+    }),
   },
   record2a: {
     id: RECORD_A_ID,
     version: 2,
     filename: "domain_category_mappings_2a.json",
+    mapping: convertDomainsToHashes({
+      "example.com": [1, 80],
+    }),
   },
   record2b: {
     id: RECORD_B_ID,
     version: 2,
     filename: "domain_category_mappings_2b.json",
+    mapping: convertDomainsToHashes({
+      "example.org": [2, 50, 4, 80],
+    }),
   },
 };
 
@@ -115,13 +145,13 @@ add_task(async function test_initial_import() {
   await promise;
 
   Assert.deepEqual(
-    SearchSERPDomainToCategoriesMap.get("example.com"),
+    await SearchSERPDomainToCategoriesMap.get("example.com"),
     [{ category: 1, score: 100 }],
     "Return value from lookup of example.com should be the same."
   );
 
   Assert.deepEqual(
-    SearchSERPDomainToCategoriesMap.get("example.org"),
+    await SearchSERPDomainToCategoriesMap.get("example.org"),
     [{ category: 2, score: 90 }],
     "Return value from lookup of example.org should be the same."
   );
@@ -167,13 +197,13 @@ add_task(async function test_update_records() {
   await promise;
 
   Assert.deepEqual(
-    SearchSERPDomainToCategoriesMap.get("example.com"),
+    await SearchSERPDomainToCategoriesMap.get("example.com"),
     [{ category: 1, score: 80 }],
     "Return value from lookup of example.com should have changed."
   );
 
   Assert.deepEqual(
-    SearchSERPDomainToCategoriesMap.get("example.org"),
+    await SearchSERPDomainToCategoriesMap.get("example.org"),
     [
       { category: 2, score: 50 },
       { category: 4, score: 80 },
@@ -224,13 +254,13 @@ add_task(async function test_delayed_initial_import() {
   await promise;
 
   Assert.deepEqual(
-    SearchSERPDomainToCategoriesMap.get("example.com"),
+    await SearchSERPDomainToCategoriesMap.get("example.com"),
     [{ category: 1, score: 100 }],
     "Return value from lookup of example.com should be the same."
   );
 
   Assert.deepEqual(
-    SearchSERPDomainToCategoriesMap.get("example.org"),
+    await SearchSERPDomainToCategoriesMap.get("example.org"),
     [{ category: 2, score: 90 }],
     "Return value from lookup of example.org should be the same."
   );
@@ -264,7 +294,7 @@ add_task(async function test_remove_record() {
   await promise;
 
   Assert.deepEqual(
-    SearchSERPDomainToCategoriesMap.get("example.com"),
+    await SearchSERPDomainToCategoriesMap.get("example.com"),
     [{ category: 1, score: 80 }],
     "Initialized properly."
   );
@@ -283,13 +313,13 @@ add_task(async function test_remove_record() {
   await promise;
 
   Assert.deepEqual(
-    SearchSERPDomainToCategoriesMap.get("example.com"),
+    await SearchSERPDomainToCategoriesMap.get("example.com"),
     [{ category: 1, score: 80 }],
     "Return value from lookup of example.com should remain unchanged."
   );
 
   Assert.deepEqual(
-    SearchSERPDomainToCategoriesMap.get("example.org"),
+    await SearchSERPDomainToCategoriesMap.get("example.org"),
     [],
     "Return value from lookup of example.org should be empty."
   );
@@ -323,7 +353,7 @@ add_task(async function test_different_versions_coexisting() {
   await promise;
 
   Assert.deepEqual(
-    SearchSERPDomainToCategoriesMap.get("example.com"),
+    await SearchSERPDomainToCategoriesMap.get("example.com"),
     [
       {
         category: 1,
@@ -334,7 +364,7 @@ add_task(async function test_different_versions_coexisting() {
   );
 
   Assert.deepEqual(
-    SearchSERPDomainToCategoriesMap.get("example.org"),
+    await SearchSERPDomainToCategoriesMap.get("example.org"),
     [
       { category: 2, score: 50 },
       { category: 4, score: 80 },
@@ -367,7 +397,7 @@ add_task(async function test_download_error() {
   await promise;
 
   Assert.deepEqual(
-    SearchSERPDomainToCategoriesMap.get("example.com"),
+    await SearchSERPDomainToCategoriesMap.get("example.com"),
     [
       {
         category: 1,
@@ -406,7 +436,7 @@ add_task(async function test_download_error() {
   await observeDownloadError;
 
   Assert.deepEqual(
-    SearchSERPDomainToCategoriesMap.get("example.com"),
+    await SearchSERPDomainToCategoriesMap.get("example.com"),
     [],
     "Domain should not exist in store."
   );