1 files changed, 18 insertions, 3 deletions
diff --git a/browser/components/search/test/browser/telemetry/browser_search_telemetry_domain_categorization_reporting_timer.js b/browser/components/search/test/browser/telemetry/browser_search_telemetry_domain_categorization_reporting_timer.js
index 9d3ac2c931..0e2d1c07fd 100644
--- a/browser/components/search/test/browser/telemetry/browser_search_telemetry_domain_categorization_reporting_timer.js
+++ b/browser/components/search/test/browser/telemetry/browser_search_telemetry_domain_categorization_reporting_timer.js
@@ -87,9 +87,20 @@ add_setup(async function () {
   await promise;
 
   registerCleanupFunction(async () => {
-    // The scheduler uses the mock idle service.
-    SearchSERPCategorizationEventScheduler.uninit();
-    SearchSERPCategorizationEventScheduler.init();
+    // Manually unload the pref so that we can check if we should wait for the
+    // the categories map to be un-initialized.
+    await SpecialPowers.popPrefEnv();
+    if (
+      !Services.prefs.getBoolPref(
+        "browser.search.serpEventTelemetryCategorization.enabled"
+      )
+    ) {
+      await waitForDomainToCategoriesUninit();
+    } else {
+      // The scheduler uses the mock idle service.
+      SearchSERPCategorizationEventScheduler.uninit();
+      SearchSERPCategorizationEventScheduler.init();
+    }
     SearchSERPTelemetry.overrideSearchTelemetryForTests();
     resetTelemetry();
   });
@@ -126,6 +137,7 @@ add_task(async function test_categorize_serp_and_wait() {
       partner_code: "ff",
       provider: "example",
       tagged: "true",
+      is_shopping_page: "false",
       num_ads_clicked: "0",
       num_ads_visible: "2",
     },
@@ -170,6 +182,7 @@ add_task(async function test_categorize_serp_open_multiple_tabs() {
       partner_code: "ff",
       provider: "example",
       tagged: "true",
+      is_shopping_page: "false",
       num_ads_clicked: "0",
       num_ads_visible: "2",
     });
@@ -223,6 +236,7 @@ add_task(async function test_categorize_serp_close_tab_and_wait() {
       partner_code: "ff",
       provider: "example",
       tagged: "true",
+      is_shopping_page: "false",
       num_ads_clicked: "0",
       num_ads_visible: "2",
     },
@@ -276,6 +290,7 @@ add_task(async function test_categorize_serp_open_ad_and_wait() {
       partner_code: "ff",
       provider: "example",
       tagged: "true",
+      is_shopping_page: "false",
       num_ads_clicked: "1",
       num_ads_visible: "2",
     },