Bug 1829412 - Simplify NimbusTestUtils.enrollmentHelper r=chumphreys,settings-reviewers,pip-reviewers,credential-management-reviewers,search-reviewers,anti-tracking-reviewers,omc-reviewers,home-newtab-reviewers,thecount,issammani,aminomancer,mconley

The enrollmentHelper was much more complicated than it needed to be. The
internal asynchrony that required awaiting an additional promise was fixed in
bug 1773583.

The returned cleanup function is no longer async, so unnecessary awaits have
been removed. This also applies to enrollWithFeatureConfig, as it is a wrapper
around enrollmentHelper.

Differential Revision: https://phabricator.services.mozilla.com/D212318

browser/base/content/test/about/browser_aboutSupport.js

@@ -107,7 +107,7 @@ add_task(async function test_nimbus_experiments() {
     }
   );
 
-  await doExperimentCleanup();
+  doExperimentCleanup();
 });
 
 add_task(async function test_remote_configuration() {
@@ -151,7 +151,7 @@ add_task(async function test_remote_configuration() {
     }
   );
 
-  await doCleanup();
+  doCleanup();
 });
 
 add_task(async function test_remote_settings() {

browser/components/aboutwelcome/tests/browser/browser_aboutwelcome_configurable_ui.js

@@ -399,7 +399,7 @@ add_task(async function test_aboutwelcome_with_url_backdrop() {
     // Expected selectors:
     [`div.outer-wrapper.onboardingContainer[style*='${TEST_BACKDROP_URL}']`]
   );
-  await doExperimentCleanup();
+  doExperimentCleanup();
   browser.closeBrowser();
 });
 
@@ -428,7 +428,7 @@ add_task(async function test_aboutwelcome_with_color_backdrop() {
     // Expected selectors:
     [`div.outer-wrapper.onboardingContainer[style*='${TEST_BACKDROP_COLOR}']`]
   );
-  await doExperimentCleanup();
+  doExperimentCleanup();
   browser.closeBrowser();
 });
 
@@ -492,7 +492,7 @@ add_task(async function test_aboutwelcome_with_text_color_override() {
     }
   );
 
-  await doExperimentCleanup();
+  doExperimentCleanup();
   await SpecialPowers.popPrefEnv();
   browser.closeBrowser();
 });
@@ -569,7 +569,7 @@ add_task(async function test_aboutwelcome_with_progress_bar() {
     );
   });
 
-  await doExperimentCleanup();
+  doExperimentCleanup();
   browser.closeBrowser();
 });
 
@@ -611,7 +611,7 @@ add_task(async function test_aboutwelcome_history_updates_disabled() {
     "No entries added to the session's history stack with history updates disabled"
   );
 
-  await doExperimentCleanup();
+  doExperimentCleanup();
   browser.closeBrowser();
 });
 
@@ -718,7 +718,7 @@ add_task(async function test_aboutwelcome_start_screen_configured() {
     ok(false, "No telemetry sent");
   }
 
-  await doExperimentCleanup();
+  doExperimentCleanup();
   browser.closeBrowser();
   sandbox.restore();
 });
@@ -743,6 +743,6 @@ add_task(async function test_aboutwelcome_rdm_property() {
     ["main.TEST_NO_RDM[no-rdm]"]
   );
 
-  await doExperimentCleanup();
+  doExperimentCleanup();
   browser.closeBrowser();
 });

browser/components/aboutwelcome/tests/browser/browser_aboutwelcome_multistage_experimentAPI.js

@@ -42,7 +42,7 @@ add_task(async function test_multistage_zeroOnboarding_experimentAPI() {
     ["div.onboardingContainer", "main.AW_STEP1"]
   );
 
-  await doExperimentCleanup();
+  doExperimentCleanup();
 });
 
 /**
@@ -256,7 +256,7 @@ add_task(async function test_multistage_aboutwelcome_experimentAPI() {
     1
   );
 
-  await doExperimentCleanup();
+  doExperimentCleanup();
 });
 
 /* Test multistage custom backdrop
@@ -307,7 +307,7 @@ add_task(async function test_multistage_aboutwelcome_backdrop() {
     [`div.outer-wrapper.onboardingContainer[style*='${TEST_BACKDROP}']`]
   );
 
-  await doExperimentCleanup();
+  doExperimentCleanup();
 });
 
 add_task(async function test_multistage_aboutwelcome_utm_term() {
@@ -380,5 +380,5 @@ add_task(async function test_multistage_aboutwelcome_utm_term() {
     BrowserTestUtils.removeTab(tab);
   });
 
-  await doExperimentCleanup();
+  doExperimentCleanup();
 });

browser/components/aboutwelcome/tests/browser/browser_aboutwelcome_multistage_transitions.js

@@ -159,7 +159,7 @@ add_task(async function test_multistage_aboutwelcome_transitions() {
     ["div.proton.transition-out .screen", "div.proton.transition- .screen-1"]
   );
 
-  await doExperimentCleanup();
+  doExperimentCleanup();
 });
 
 /**
@@ -215,5 +215,5 @@ add_task(async function test_multistage_aboutwelcome_transitions_off() {
     ["div.proton.transition-out .screen-0"]
   );
 
-  await doExperimentCleanup();
+  doExperimentCleanup();
 });

browser/components/aboutwelcome/tests/browser/browser_aboutwelcome_rtamo.js

@@ -239,7 +239,7 @@ add_task(async function test_rtamo_over_experiments() {
     []
   );
 
-  await doExperimentCleanup();
+  doExperimentCleanup();
 
   browser = await openRTAMOWelcomePage();
 

browser/components/newtab/test/browser/abouthomecache/browser_experiments_api_control.js

@@ -42,7 +42,7 @@ add_task(async function test_experiments_api_control() {
       AboutHomeStartupCache.CACHE_RESULT_SCALARS.DISABLED
     );
 
-    await doEnrollmentCleanup();
+    doEnrollmentCleanup();
   });
 
   // Now the enabled case.
@@ -62,6 +62,6 @@ add_task(async function test_experiments_api_control() {
 
     await simulateRestart(browser);
     await ensureCachedAboutHome(browser);
-    await doEnrollmentCleanup();
+    doEnrollmentCleanup();
   });
 });

browser/components/newtab/test/browser/browser_newtab_ping.js

@@ -191,7 +191,7 @@ add_task(async function test_newtab_doesnt_send_nimbus() {
     ).length;
   }, "Waiting for sessions to clean up.");
   // Session ended without a ping being sent. Success!
-  await doEnrollmentCleanup();
+  doEnrollmentCleanup();
   await SpecialPowers.popPrefEnv();
 });
 

browser/components/preferences/tests/browser_windows_launch_on_login.js

@@ -68,7 +68,7 @@ add_task(async function test_check_uncheck_checkbox() {
 
     gBrowser.removeCurrentTab();
   });
-  await doCleanup();
+  doCleanup();
 });
 
 add_task(async function create_external_regkey() {
@@ -111,7 +111,7 @@ add_task(async function create_external_regkey() {
 
     gBrowser.removeCurrentTab();
   });
-  await doCleanup();
+  doCleanup();
 });
 
 add_task(async function delete_external_regkey() {
@@ -142,7 +142,7 @@ add_task(async function delete_external_regkey() {
 
     gBrowser.removeCurrentTab();
   });
-  await doCleanup();
+  doCleanup();
 });
 
 registerCleanupFunction(async function () {

browser/components/preferences/tests/browser_windows_launch_on_login_msix.js

@@ -52,7 +52,7 @@ add_task(async function test_check_uncheck_checkbox() {
   }, "Wait for async get enabled operation to return false");
 
   gBrowser.removeCurrentTab();
-  await doCleanup();
+  doCleanup();
 });
 
 add_task(async function enable_external_startuptask() {
@@ -76,7 +76,7 @@ add_task(async function enable_external_startuptask() {
   ok(launchOnLoginCheckbox.checked, "Autostart checkbox automatically checked");
 
   gBrowser.removeCurrentTab();
-  await doCleanup();
+  doCleanup();
 });
 
 add_task(async function disable_external_startuptask() {
@@ -101,5 +101,5 @@ add_task(async function disable_external_startuptask() {
   );
 
   gBrowser.removeCurrentTab();
-  await doCleanup();
+  doCleanup();
 });

browser/components/privatebrowsing/test/browser/browser_privatebrowsing_about_nimbus.js

@@ -67,7 +67,7 @@ add_task(async function test_experiment_plain_text() {
   });
 
   await BrowserTestUtils.closeWindow(win);
-  await doExperimentCleanup();
+  doExperimentCleanup();
 });
 
 add_task(async function test_experiment_info_disabled() {
@@ -93,7 +93,7 @@ add_task(async function test_experiment_info_disabled() {
   });
 
   await BrowserTestUtils.closeWindow(win);
-  await doExperimentCleanup();
+  doExperimentCleanup();
 });
 
 add_task(async function test_experiment_promo_disabled() {
@@ -120,7 +120,7 @@ add_task(async function test_experiment_promo_disabled() {
   });
 
   await BrowserTestUtils.closeWindow(win);
-  await doExperimentCleanup();
+  doExperimentCleanup();
 });
 
 add_task(async function test_experiment_format_urls() {
@@ -164,7 +164,7 @@ add_task(async function test_experiment_format_urls() {
   });
 
   await BrowserTestUtils.closeWindow(win);
-  await doExperimentCleanup();
+  doExperimentCleanup();
 });
 
 add_task(async function test_experiment_click_info_telemetry() {
@@ -199,7 +199,7 @@ add_task(async function test_experiment_click_info_telemetry() {
   );
 
   await BrowserTestUtils.closeWindow(win);
-  await doExperimentCleanup();
+  doExperimentCleanup();
 });
 
 add_task(async function test_experiment_click_promo_telemetry() {
@@ -250,7 +250,7 @@ add_task(async function test_experiment_click_promo_telemetry() {
   );
 
   await BrowserTestUtils.closeWindow(win);
-  await doExperimentCleanup();
+  doExperimentCleanup();
 });
 
 add_task(async function test_experiment_bottom_promo() {
@@ -318,7 +318,7 @@ add_task(async function test_experiment_bottom_promo() {
   });
 
   await BrowserTestUtils.closeWindow(win);
-  await doExperimentCleanup();
+  doExperimentCleanup();
 });
 
 add_task(async function test_experiment_below_search_promo() {
@@ -390,7 +390,7 @@ add_task(async function test_experiment_below_search_promo() {
   });
 
   await BrowserTestUtils.closeWindow(win);
-  await doExperimentCleanup();
+  doExperimentCleanup();
 });
 
 add_task(async function test_experiment_top_promo() {
@@ -455,5 +455,5 @@ add_task(async function test_experiment_top_promo() {
   });
 
   await BrowserTestUtils.closeWindow(win);
-  await doExperimentCleanup();
+  doExperimentCleanup();
 });

browser/components/privatebrowsing/test/browser/browser_privatebrowsing_about_nimbus_dismiss.js

@@ -65,7 +65,7 @@ add_task(async function test_experiment_messaging_system_dismiss() {
 
   await BrowserTestUtils.closeWindow(win1);
   await BrowserTestUtils.closeWindow(win2);
-  await doExperimentCleanup();
+  doExperimentCleanup();
 });
 
 add_task(async function test_experiment_messaging_show_default_on_dismiss() {
@@ -135,5 +135,5 @@ add_task(async function test_experiment_messaging_show_default_on_dismiss() {
 
   await BrowserTestUtils.closeWindow(win1);
   await BrowserTestUtils.closeWindow(win2);
-  await doExperimentCleanup();
+  doExperimentCleanup();
 });

browser/components/privatebrowsing/test/browser/browser_privatebrowsing_about_nimbus_impressions.js

@@ -122,5 +122,5 @@ add_task(async function test_experiment_messaging_system_impressions() {
   await BrowserTestUtils.closeWindow(win1);
   await BrowserTestUtils.closeWindow(win2);
   await BrowserTestUtils.closeWindow(win3);
-  await doExperimentCleanup();
+  doExperimentCleanup();
 });

browser/components/privatebrowsing/test/browser/browser_privatebrowsing_about_nimbus_messaging.js

@@ -66,7 +66,7 @@ add_task(async function test_experiment_messaging_system() {
   });
 
   await BrowserTestUtils.closeWindow(win);
-  await doExperimentCleanup();
+  doExperimentCleanup();
 });
 
 add_task(async function test_experiment_promo_action() {
@@ -156,7 +156,7 @@ add_task(async function test_experiment_promo_action() {
     "Should be called with right URL"
   );
 
-  await doExperimentCleanup();
+  doExperimentCleanup();
 });
 
 add_task(async function test_experiment_open_spotlight_action() {
@@ -243,5 +243,5 @@ add_task(async function test_experiment_open_spotlight_action() {
     "Should be called with metrics property set as allow for experiments"
   );
 
-  await doExperimentCleanup();
+  doExperimentCleanup();
 });

browser/components/search/test/browser/browser_defaultPrivate_nimbus.js

@@ -131,7 +131,7 @@ add_task(async function test_nimbus_experiment() {
   );
   reloadObserved =
     SearchTestUtils.promiseSearchNotification("engines-reloaded");
-  await doExperimentCleanup();
+  doExperimentCleanup();
   await reloadObserved;
   Assert.equal(
     Services.search.defaultPrivateEngine.name,
@@ -167,7 +167,7 @@ add_task(async function test_nimbus_experiment_urlbar_result_enabled() {
   );
   reloadObserved =
     SearchTestUtils.promiseSearchNotification("engines-reloaded");
-  await doExperimentCleanup();
+  doExperimentCleanup();
   await reloadObserved;
   Assert.equal(
     Services.search.defaultPrivateEngine.name,
@@ -193,6 +193,6 @@ add_task(async function test_non_experiment_prefs() {
     },
   });
   Assert.equal(uiPref(), false, "Pref did not change without experiment");
-  await doExperimentCleanup();
+  doExperimentCleanup();
   await SpecialPowers.popPrefEnv();
 });

browser/components/search/test/browser/browser_search_nimbus_reload.js

@@ -46,7 +46,7 @@ add_task(async function test_engines_reloaded_nimbus() {
     "Should have expected value"
   );
 
-  await doExperimentCleanup();
+  doExperimentCleanup();
 
   Assert.equal(reloadSpy.callCount, 2, "Called by experiment unenrollment");
 

browser/components/search/test/browser/telemetry/browser_search_glean_serp_event_telemetry_categorization_enabled_by_nimbus_variable.js

@@ -170,7 +170,7 @@ add_task(async function test_enable_experiment_when_pref_is_not_enabled() {
   resetTelemetry();
 
   info("End experiment.");
-  await doExperimentCleanup();
+  doExperimentCleanup();
   await waitForDomainToCategoriesUninit();
 
   Assert.equal(

browser/components/shell/test/browser_doesAppNeedPin.js

@@ -38,7 +38,7 @@ add_task(async function remote_disable() {
     "Pinning disabled via nimbus"
   );
 
-  await doCleanup();
+  doCleanup();
 });
 
 add_task(async function restore_default() {

browser/components/shell/test/browser_setDefaultBrowser.js

@@ -80,7 +80,7 @@ add_task(async function remote_disable() {
   );
   Assert.ok(setDefaultStub.called, "Used plain set default insteead");
 
-  await doCleanup();
+  doCleanup();
 });
 
 add_task(async function restore_default() {
@@ -144,5 +144,5 @@ add_task(async function ensure_fallback() {
   );
   Assert.ok(setDefaultStub.called, "Fallbacked to plain set default");
 
-  await doCleanup();
+  doCleanup();
 });

browser/components/shell/test/browser_setDefaultPDFHandler.js

@@ -87,7 +87,7 @@ add_task(async function remoteEnableWithPDF() {
     [".pdf", "FirefoxPDF"],
   ]);
 
-  await doCleanup();
+  doCleanup();
 });
 
 add_task(async function remoteEnableWithPDF_testOnlyReplaceBrowsers() {
@@ -149,7 +149,7 @@ add_task(async function remoteEnableWithPDF_testOnlyReplaceBrowsers() {
     `Will not take default from non-browser`
   );
 
-  await doCleanup();
+  doCleanup();
 });
 
 add_task(async function remoteEnableWithoutPDF() {
@@ -181,7 +181,7 @@ add_task(async function remoteEnableWithoutPDF() {
   Assert.ok(setDefaultBrowserUserChoiceStub.called);
   Assert.deepEqual(setDefaultBrowserUserChoiceStub.firstCall.args, [aumi, []]);
 
-  await doCleanup();
+  doCleanup();
 });
 
 add_task(async function remoteDisable() {
@@ -212,7 +212,7 @@ add_task(async function remoteDisable() {
   Assert.ok(setDefaultBrowserUserChoiceStub.notCalled);
   Assert.ok(setDefaultStub.called);
 
-  await doCleanup();
+  doCleanup();
 });
 
 add_task(async function test_setAsDefaultPDFHandler_knownBrowser() {

browser/components/tests/browser/browser_browserGlue_upgradeDialog_trigger.js

@@ -68,7 +68,7 @@ add_task(async function remote_disabled() {
     "disabled",
   ]);
 
-  await doCleanup();
+  doCleanup();
   await cleanupUpgrade();
 });
 

browser/components/tests/browser/browser_default_webprotocol_handler_mailto.js

@@ -244,7 +244,7 @@ add_task(async function check_no_button() {
     );
   });
 
-  await cleanup();
+  cleanup();
 });
 
 add_task(async function check_x_button() {
@@ -283,7 +283,7 @@ add_task(async function check_x_button() {
     );
   });
 
-  await cleanup();
+  cleanup();
 });
 
 add_task(async function check_x_button() {
@@ -297,7 +297,7 @@ add_task(async function check_x_button() {
     );
   });
 
-  await cleanup();
+  cleanup();
 });
 
 add_task(async function check_bar_is_not_shown() {
@@ -311,5 +311,5 @@ add_task(async function check_bar_is_not_shown() {
     );
   });
 
-  await cleanup();
+  cleanup();
 });

browser/components/urlbar/tests/UrlbarTestUtils.sys.mjs

@@ -1292,7 +1292,7 @@ export var UrlbarTestUtils = {
       // If `doCleanup()` has already been called (i.e., by the caller), it will
       // throw an error here.
       try {
-        await doCleanup();
+        doCleanup();
       } catch (error) {}
     });
 

browser/components/urlbar/tests/quicksuggest/QuickSuggestTestUtils.sys.mjs

@@ -820,7 +820,7 @@ class _QuickSuggestTestUtils {
 
     return async () => {
       this.#log("enrollExperiment.cleanup", "Awaiting experiment cleanup");
-      await doExperimentCleanup();
+      doExperimentCleanup();
 
       // The same pref updates will be triggered by unenrollment, so wait for
       // them again.

browser/components/urlbar/tests/quicksuggest/unit/test_merinoClient.js

@@ -635,5 +635,5 @@ async function withExperiment(values, callback) {
     }
   );
   await callback();
-  await doExperimentCleanup();
+  doExperimentCleanup();
 }

browser/components/urlbar/tests/quicksuggest/unit/test_quicksuggest_migrate_v2.js

@@ -1343,5 +1343,5 @@ async function withOnlineExperiment(callback) {
     }
   );
   await callback();
-  await doExperimentCleanup();
+  doExperimentCleanup();
 }

browser/extensions/formautofill/test/browser/address/browser_address_autofill_nimbus.js

@@ -41,7 +41,7 @@ add_task(async function test_address_autofill_feature_enabled() {
     "Address autofill should be available when feature is enabled in nimbus."
   );
 
-  await cleanupExperiment();
+  cleanupExperiment();
 });
 
 add_task(async function test_address_autofill_feature_disabled() {
@@ -66,5 +66,5 @@ add_task(async function test_address_autofill_feature_disabled() {
     "Address autofill shouldn't be available when feature is off in nimbus."
   );
 
-  await cleanupExperiment();
+  cleanupExperiment();
 });

toolkit/components/antitracking/test/browser/browser_urlQueryStringStripping_nimbus.js

@@ -133,7 +133,7 @@ add_task(async function test() {
         );
       }
 
-      await doExperimentCleanup();
+      doExperimentCleanup();
     }
   }
 

toolkit/components/contentrelevancy/tests/browser/browser_contentrelevancy_nimbus.js

@@ -55,7 +55,7 @@ add_task(async function test_NimbusIntegration_enable() {
     "The timer callback should be called"
   );
 
-  await doExperimentCleanup();
+  doExperimentCleanup();
   gSandbox.restore();
 });
 
@@ -88,6 +88,6 @@ add_task(async function test_NimbusIntegration_disable() {
     "The timer callback should not be called"
   );
 
-  await doExperimentCleanup();
+  doExperimentCleanup();
   gSandbox.restore();
 });

toolkit/components/nimbus/test/NimbusTestUtils.sys.mjs

@@ -233,22 +233,23 @@ export const ExperimentFakes = {
       ],
       isRollout,
     });
-    let { enrollmentPromise, doExperimentCleanup } = this.enrollmentHelper(
-      recipe,
-      { manager, source }
-    );
+    const doEnrollmentCleanup = await this.enrollmentHelper(recipe, {
+      manager,
+      source,
+    });
 
-    await enrollmentPromise;
-
-    return doExperimentCleanup;
+    return doEnrollmentCleanup;
   },
   /**
-   * Enroll in the given recipe.
+   * Attempt enrollment in the given recipe.
    *
-   * NB: It is unnecessary to await the enrollmentPromise.
-   *     See bug 1773583 and bug 1829412.
+   * This will evaluate bucketing, so it is possible that enrollment will *not*
+   * occur.
+   *
+   * If you are testing a feature, you likely want to use enrollInFeatureConfig,
+   * which will guarantee successful enrollment.
    */
-  enrollmentHelper(
+  async enrollmentHelper(
     recipe,
     { manager = lazy.ExperimentAPI._manager, source = "enrollmentHelper" } = {}
   ) {
@@ -256,38 +257,15 @@ export const ExperimentFakes = {
       throw new Error("Enrollment helper expects a recipe");
     }
 
-    let enrollmentPromise = new Promise(resolve =>
-      manager.store.on(`update:${recipe.slug}`, (event, experiment) => {
-        if (experiment.active) {
+    const enrollment = await manager.enroll(recipe, source);
+    if (enrollment?.active) {
       manager.store._syncToChildren({ flush: true });
-          resolve(experiment);
     }
-      })
-    );
-    let unenrollCompleted = slug =>
-      new Promise(resolve =>
-        manager.store.on(`update:${slug}`, (event, experiment) => {
-          if (!experiment.active) {
-            // Removes recipe from file storage which
-            // (normally the users archive of past experiments)
-            manager.store._deleteForTests(recipe.slug);
-            resolve();
-          }
-        })
-      );
-    let doExperimentCleanup = async () => {
-      const experiment = manager.store.get(recipe.slug);
-      let promise = unenrollCompleted(experiment.slug);
-      manager.unenroll(experiment.slug, "cleanup");
-      await promise;
+
+    return function doEnrollmentCleanup() {
+      manager.unenroll(enrollment.slug, "cleanup");
+      manager.store._deleteForTests(enrollment.slug);
     };
-
-    if (!manager.store._isReady) {
-      throw new Error("Manager store not ready, call `manager.onStartup`");
-    }
-    manager.enroll(recipe, source);
-
-    return { enrollmentPromise, doExperimentCleanup };
   },
   async cleanupAll(slugs, { manager = lazy.ExperimentAPI._manager } = {}) {
     function unenrollCompleted(slug) {

toolkit/components/nimbus/test/browser/browser_experiment_evaluate_jexl.js

@@ -80,10 +80,7 @@ add_task(async function test_evaluate_active_experiments_activeExperiments() {
   recipe.branches[0].slug = "mochitest-active-foo";
   delete recipe.branches[1];
 
-  let { enrollmentPromise, doExperimentCleanup } =
-    ExperimentFakes.enrollmentHelper(recipe);
-
-  await enrollmentPromise;
+  const doExperimentCleanup = await ExperimentFakes.enrollmentHelper(recipe);
 
   Assert.equal(
     await CONTEXT.evaluateJexl(`"${slug}" in activeExperiments`, FAKE_CONTEXT),
@@ -100,5 +97,5 @@ add_task(async function test_evaluate_active_experiments_activeExperiments() {
     "should not find an experiment that doesn't exist"
   );
 
-  await doExperimentCleanup();
+  doExperimentCleanup();
 });

toolkit/components/nimbus/test/browser/browser_experiment_single_feature_enrollment.js

@@ -76,12 +76,12 @@ const SINGLE_FEATURE_RECIPE = {
 const SYNC_DATA_PREF_BRANCH = "nimbus.syncdatastore.";
 
 add_task(async function test_TODO() {
-  let { enrollmentPromise, doExperimentCleanup } =
-    ExperimentFakes.enrollmentHelper(SINGLE_FEATURE_RECIPE);
   let sandbox = sinon.createSandbox();
   let stub = sandbox.stub(ExperimentAPI, "recordExposureEvent");
 
-  await enrollmentPromise;
+  const doExperimentCleanup = await ExperimentFakes.enrollmentHelper(
+    SINGLE_FEATURE_RECIPE
+  );
 
   Assert.ok(
     ExperimentAPI.getExperiment({ featureId: "urlbar" }),
@@ -122,7 +122,7 @@ add_task(async function test_TODO() {
     "Should have expected featureId"
   );
 
-  await doExperimentCleanup();
+  doExperimentCleanup();
   sandbox.restore();
   NimbusFeatures.urlbar._didSendExposureEvent = false;
 });

toolkit/components/nimbus/test/browser/browser_nimbus_telemetry.js

@@ -59,7 +59,7 @@ add_task(async function test_experiment_enroll_unenroll_Telemetry() {
     EVENT_FILTER
   );
 
-  await cleanup();
+  cleanup();
 
   TelemetryTestUtils.assertEvents(
     [
@@ -112,7 +112,7 @@ add_task(async function test_experiment_expose_Telemetry() {
     EVENT_FILTER
   );
 
-  await cleanup();
+  cleanup();
 });
 
 add_task(async function test_rollout_expose_Telemetry() {
@@ -154,5 +154,5 @@ add_task(async function test_rollout_expose_Telemetry() {
     EVENT_FILTER
   );
 
-  await cleanup();
+  cleanup();
 });

toolkit/components/nimbus/test/browser/browser_prefs.js

@@ -65,7 +65,7 @@ add_task(async function test_prefs_priority() {
     );
 
     for (const cleanup of enrollmentCleanup) {
-      await cleanup();
+      cleanup();
     }
 
     Services.prefs.deleteBranch(pref);

toolkit/components/nimbus/test/browser/browser_remotesettingsexperimentloader_remote_defaults.js

@@ -345,15 +345,7 @@ add_task(async function test_finalizeRemoteConfigs_cleanup() {
       source: "rs-loader",
     }
   );
-  let stubFoo = sinon.stub().callsFake((...args) => {
-    dump(`
-!!! stubfoo called with ${JSON.stringify(args)}
-
-
-${new Error().stack}
-!!!
-`);
-  });
+  let stubFoo = sinon.stub();
   let stubBar = sinon.stub();
   featureFoo.onUpdate(stubFoo);
   featureBar.onUpdate(stubBar);
@@ -417,7 +409,7 @@ ${new Error().stack}
     "Pref was cleared"
   );
 
-  await fooCleanup();
+  fooCleanup();
   // This will also remove the inactive recipe from the store
   // the previous update (from recipe not seen code path)
   // only sets the recipe as inactive
@@ -586,7 +578,7 @@ add_task(async function remote_defaults_variables_storage() {
     "Test types are returned correctly"
   );
 
-  await doCleanup();
+  doCleanup();
 
   Assert.equal(
     Services.prefs.getIntPref(`${SYNC_DEFAULTS_PREF_BRANCH}bar.storage`, -1),

toolkit/components/nimbus/test/unit/test_ExperimentAPI_ExperimentFeature.js

@@ -228,7 +228,7 @@ add_task(async function test_allow_multiple_exposure_events() {
   Assert.equal(3, exposureEvents.length);
 
   sandbox.restore();
-  await doExperimentCleanup();
+  doExperimentCleanup();
 });
 
 add_task(async function test_onUpdate_before_store_ready() {

toolkit/components/nimbus/test/unit/test_ExperimentAPI_ExperimentFeature_getVariable.js

@@ -144,7 +144,7 @@ add_task(async function test_ExperimentFeature_getVariable_precedence() {
 
   // Cleanup
   Services.prefs.deleteBranch(TEST_PREF_BRANCH);
-  await doExperimentCleanup();
+  doExperimentCleanup();
   sandbox.restore();
 });
 
@@ -191,6 +191,6 @@ add_task(async function test_ExperimentFeature_getVariable_partial_values() {
   // Cleanup
   Services.prefs.getDefaultBranch("").deleteBranch(TEST_PREF_BRANCH);
   Services.prefs.deleteBranch(TEST_PREF_BRANCH);
-  await doExperimentCleanup();
+  doExperimentCleanup();
   sandbox.restore();
 });

toolkit/components/nimbus/test/unit/test_ExperimentManager_enroll.js

@@ -534,10 +534,10 @@ add_task(async function test_rollout_experiment_no_conflict() {
 
   await ExperimentFakes.enrollmentHelper(experiment, {
     manager,
-  }).enrollmentPromise;
+  });
   await ExperimentFakes.enrollmentHelper(rollout, {
     manager,
-  }).enrollmentPromise;
+  });
 
   Assert.ok(
     manager.store.get(experiment.slug).active,
@@ -854,10 +854,9 @@ add_task(async function test_featureIds_is_stored() {
 
   await manager.onStartup();
 
-  const { enrollmentPromise, doExperimentCleanup } =
-    ExperimentFakes.enrollmentHelper(recipe, { manager });
-
-  await enrollmentPromise;
+  const doExperimentCleanup = await ExperimentFakes.enrollmentHelper(recipe, {
+    manager,
+  });
 
   Assert.ok(manager.store.addEnrollment.calledOnce, "experiment is stored");
   let [enrollment] = manager.store.addEnrollment.firstCall.args;
@@ -868,7 +867,7 @@ add_task(async function test_featureIds_is_stored() {
     "Has expected value"
   );
 
-  await doExperimentCleanup();
+  doExperimentCleanup();
 
   await assertEmptyStore(manager.store);
 });
@@ -879,7 +878,7 @@ add_task(async function experiment_and_rollout_enroll_and_cleanup() {
 
   await manager.onStartup();
 
-  let rolloutCleanup = await ExperimentFakes.enrollWithFeatureConfig(
+  let doRolloutCleanup = await ExperimentFakes.enrollWithFeatureConfig(
     {
       featureId: "aboutwelcome",
       value: { enabled: true },
@@ -890,7 +889,7 @@ add_task(async function experiment_and_rollout_enroll_and_cleanup() {
     }
   );
 
-  let experimentCleanup = await ExperimentFakes.enrollWithFeatureConfig(
+  let doExperimentCleanup = await ExperimentFakes.enrollWithFeatureConfig(
     {
       featureId: "aboutwelcome",
       value: { enabled: true },
@@ -907,7 +906,7 @@ add_task(async function experiment_and_rollout_enroll_and_cleanup() {
     )
   );
 
-  await experimentCleanup();
+  doExperimentCleanup();
 
   Assert.ok(
     !Services.prefs.getBoolPref(
@@ -921,7 +920,7 @@ add_task(async function experiment_and_rollout_enroll_and_cleanup() {
     )
   );
 
-  await rolloutCleanup();
+  doRolloutCleanup();
 
   Assert.ok(
     !Services.prefs.getBoolPref(

toolkit/components/nimbus/test/unit/test_ExperimentManager_prefs.js

@@ -2698,7 +2698,7 @@ add_task(async function test_prefChanged_noPrefSet() {
 
             assertNoObservers(manager);
 
-            await doEnrollmentCleanup();
+            doEnrollmentCleanup();
             await assertEmptyStore(store);
 
             Services.prefs.deleteBranch(pref);

toolkit/components/nimbus/test/unit/test_ExperimentManager_unenroll.js

@@ -191,7 +191,7 @@ add_task(async function test_unenroll_uploadPref() {
   const recipe = ExperimentFakes.recipe("foo");
 
   await manager.onStartup();
-  await ExperimentFakes.enrollmentHelper(recipe, { manager }).enrollmentPromise;
+  await ExperimentFakes.enrollmentHelper(recipe, { manager });
 
   Assert.equal(
     manager.store.get(recipe.slug).active,

toolkit/components/nimbus/test/unit/test_NimbusTestUtils.js

@@ -32,10 +32,9 @@ add_task(async function test_enrollmentHelper() {
 
   await manager.onStartup();
 
-  let { enrollmentPromise, doExperimentCleanup } =
-    ExperimentFakes.enrollmentHelper(recipe, { manager });
-
-  await enrollmentPromise;
+  const doEnrollmentCleanup = await ExperimentFakes.enrollmentHelper(recipe, {
+    manager,
+  });
 
   Assert.ok(manager.store.getAllActiveExperiments().length === 1, "Enrolled");
   Assert.equal(
@@ -48,7 +47,7 @@ add_task(async function test_enrollmentHelper() {
     "Sync pref cache set"
   );
 
-  await doExperimentCleanup();
+  doEnrollmentCleanup();
 
   Assert.ok(manager.store.getAll().length === 0, "Cleanup done");
   Assert.ok(
@@ -60,7 +59,7 @@ add_task(async function test_enrollmentHelper() {
 add_task(async function test_enrollWithFeatureConfig() {
   let manager = ExperimentFakes.manager();
   await manager.onStartup();
-  let doExperimentCleanup = await ExperimentFakes.enrollWithFeatureConfig(
+  let doEnrollmentCleanup = await ExperimentFakes.enrollWithFeatureConfig(
     {
       featureId: "enrollWithFeatureConfig",
       value: { enabled: true },
@@ -73,7 +72,7 @@ add_task(async function test_enrollWithFeatureConfig() {
     "Enrolled successfully"
   );
 
-  await doExperimentCleanup();
+  doEnrollmentCleanup();
 
   Assert.ok(
     !manager.store.hasExperimentForFeature("enrollWithFeatureConfig"),

toolkit/components/nimbus/test/unit/test_localization.js

@@ -230,9 +230,9 @@ add_task(async function test_getLocalizedValue() {
     localizations: LOCALIZATIONS,
   });
 
-  const { enrollmentPromise, doExperimentCleanup } =
-    ExperimentFakes.enrollmentHelper(experiment);
-  await enrollmentPromise;
+  const doExperimentCleanup = await ExperimentFakes.enrollmentHelper(
+    experiment
+  );
 
   const enrollment = manager.store.getExperimentForFeature(FEATURE_ID);
 
@@ -254,7 +254,7 @@ add_task(async function test_getLocalizedValue() {
     "_getLocalizedValue() with a nested localization"
   );
 
-  await doExperimentCleanup();
+  doExperimentCleanup();
   await cleanupStore(manager.store);
   sandbox.reset();
 });
@@ -294,7 +294,7 @@ add_task(async function test_getLocalizedValue_unenroll_missingEntry() {
     localizations: LOCALIZATIONS,
   });
 
-  await ExperimentFakes.enrollmentHelper(experiment).enrollmentPromise;
+  await ExperimentFakes.enrollmentHelper(experiment);
 
   const enrollment = manager.store.getExperimentForFeature(FEATURE_ID);
 
@@ -378,7 +378,7 @@ add_task(async function test_getLocalizedValue_unenroll_missingEntry() {
     },
   });
 
-  await ExperimentFakes.enrollmentHelper(experiment).enrollmentPromise;
+  await ExperimentFakes.enrollmentHelper(experiment);
 
   const enrollment = manager.store.getExperimentForFeature(FEATURE_ID);
 
@@ -450,9 +450,9 @@ add_task(async function test_getVariables() {
     localizations: LOCALIZATIONS,
   });
 
-  const { enrollmentPromise, doExperimentCleanup } =
-    ExperimentFakes.enrollmentHelper(experiment);
-  await enrollmentPromise;
+  const doExperimentCleanup = await ExperimentFakes.enrollmentHelper(
+    experiment
+  );
 
   Assert.deepEqual(
     FEATURE.getAllVariables(),
@@ -484,7 +484,7 @@ add_task(async function test_getVariables() {
     "getVariable() returns substitutions inside arrays"
   );
 
-  await doExperimentCleanup();
+  doExperimentCleanup();
   await cleanupStore(manager.store);
   sandbox.reset();
 });
@@ -584,13 +584,7 @@ add_task(async function test_getVariables_fallback() {
   );
 
   // Enroll in the rollout.
-  {
-    const { enrollmentPromise, doExperimentCleanup } =
-      ExperimentFakes.enrollmentHelper(recipes.rollout);
-    await enrollmentPromise;
-
-    cleanup.rollout = doExperimentCleanup;
-  }
+  cleanup.rollout = await ExperimentFakes.enrollmentHelper(recipes.rollout);
 
   Assert.deepEqual(
     FEATURE.getAllVariables({ defaultValues: { waldo: ["default-value"] } }),
@@ -620,13 +614,9 @@ add_task(async function test_getVariables_fallback() {
   );
 
   // Enroll in the experiment.
-  {
-    const { enrollmentPromise, doExperimentCleanup } =
-      ExperimentFakes.enrollmentHelper(recipes.experiment);
-    await enrollmentPromise;
-
-    cleanup.experiment = doExperimentCleanup;
-  }
+  cleanup.experiment = await ExperimentFakes.enrollmentHelper(
+    recipes.experiment
+  );
 
   Assert.deepEqual(
     FEATURE.getAllVariables({ defaultValues: { waldo: ["default-value"] } }),
@@ -791,7 +781,7 @@ add_task(async function test_getVariables_fallback_unenroll() {
   ];
 
   for (const recipe of recipes) {
-    await ExperimentFakes.enrollmentHelper(recipe).enrollmentPromise;
+    await ExperimentFakes.enrollmentHelper(recipe);
   }
 
   Assert.deepEqual(FEATURE.getAllVariables(), {
@@ -1140,9 +1130,7 @@ add_task(async function test_updateRecipes_unenroll_missingEntry() {
   await manager.onStartup();
   await manager.store.ready();
 
-  await ExperimentFakes.enrollmentHelper(recipe, {
-    source: "rs-loader",
-  }).enrollmentPromise;
+  await ExperimentFakes.enrollmentHelper(recipe, { source: "rs-loader" });
   Assert.ok(
     !!manager.store.getExperimentForFeature(FEATURE_ID),
     "Should be enrolled in the experiment"
@@ -1295,9 +1283,7 @@ add_task(async function test_updateRecipes_unenroll_missingLocale() {
   await manager.onStartup();
   await manager.store.ready();
 
-  await ExperimentFakes.enrollmentHelper(recipe, {
-    source: "rs-loader",
-  }).enrollmentPromise;
+  await ExperimentFakes.enrollmentHelper(recipe, { source: "rs-loader" });
   Assert.ok(
     !!manager.store.getExperimentForFeature(FEATURE_ID),
     "Should be enrolled in the experiment"

toolkit/components/normandy/test/browser/browser_LegacyHeartbeat.js

@@ -139,7 +139,7 @@ decorate_task(
   async function testLegacyHeartbeatTrigger({ heartbeatClassStub }) {
     const sandbox = sinon.createSandbox();
 
-    const cleanupEnrollment = await ExperimentFakes.enrollWithFeatureConfig({
+    const doEnrollmentCleanup = await ExperimentFakes.enrollWithFeatureConfig({
       featureId: "legacyHeartbeat",
       value: {
         survey: SURVEY,
@@ -158,7 +158,7 @@ decorate_task(
       );
       assertSurvey(heartbeatClassStub.args[0][1], SURVEY);
 
-      await cleanupEnrollment();
+      doEnrollmentCleanup();
     } finally {
       sandbox.restore();
     }
@@ -170,7 +170,7 @@ decorate_task(
   async function testLegacyHeartbeatPingPayload() {
     const sandbox = sinon.createSandbox();
 
-    const cleanupEnrollment = await ExperimentFakes.enrollWithFeatureConfig({
+    const doEnrollmentCleanup = await ExperimentFakes.enrollWithFeatureConfig({
       featureId: "legacyHeartbeat",
       value: {
         survey: SURVEY,
@@ -215,7 +215,7 @@ decorate_task(
       Assert.ok(result.valid);
       Assert.equal(payload.surveyVersion, "1");
 
-      await cleanupEnrollment();
+      doEnrollmentCleanup();
     } finally {
       ShowHeartbeatAction.overrideHeartbeatForTests();
       sandbox.restore();

toolkit/components/passwordmgr/test/browser/browser_autocomplete_import.js

@@ -46,8 +46,8 @@ add_setup(async function setup() {
   // This makes the last autocomplete test *not* show import suggestions.
   Services.prefs.setIntPref("signon.suggestImportCount", 3);
 
-  registerCleanupFunction(async () => {
-    await doExperimentCleanup();
+  registerCleanupFunction(() => {
+    doExperimentCleanup();
     debounce.restore();
     importable.restore();
     migrator.restore();

toolkit/components/pictureinpicture/tests/browser_nimbusDisplayDuration.js

@@ -83,7 +83,7 @@ add_task(async function test_experiment_displayDuration_end_date_was_reached() {
     }
   );
 
-  await doExperimentCleanup();
+  doExperimentCleanup();
 });
 
 /**
@@ -131,7 +131,7 @@ add_task(async function test_experiment_displayDuration_end_date_not_reached() {
     }
   );
 
-  await doExperimentCleanup();
+  doExperimentCleanup();
 });
 
 /**
@@ -172,7 +172,7 @@ add_task(async function test_experiment_displayDuration_negative_duration() {
     }
   );
 
-  await doExperimentCleanup();
+  doExperimentCleanup();
 });
 
 /**
@@ -210,5 +210,5 @@ add_task(async function test_experiment_displayDuration_already_icon() {
     }
   );
 
-  await doExperimentCleanup();
+  doExperimentCleanup();
 });

toolkit/components/pictureinpicture/tests/browser_nimbusFirstTimeStyleVariant.js

@@ -69,10 +69,6 @@ add_task(async function test_experiment_toggle_style() {
     },
   });
 
-  registerCleanupFunction(async function () {
-    await doExperimentCleanup();
-  });
-
   await BrowserTestUtils.withNewTab(
     {
       gBrowser,
@@ -115,4 +111,6 @@ add_task(async function test_experiment_toggle_style() {
       );
     }
   );
+
+  doExperimentCleanup();
 });

toolkit/components/pictureinpicture/tests/browser_nimbusMessageFirstTimePip.js

@@ -72,10 +72,6 @@ add_task(async function test_experiment_message() {
     },
   });
 
-  registerCleanupFunction(async function () {
-    await doExperimentCleanup();
-  });
-
   await BrowserTestUtils.withNewTab(
     {
       gBrowser,
@@ -118,4 +114,6 @@ add_task(async function test_experiment_message() {
       );
     }
   );
+
+  doExperimentCleanup();
 });

toolkit/components/pictureinpicture/tests/browser_nimbusShowIconOnly.js

@@ -71,10 +71,6 @@ add_task(async function test_experiment_iconOnly() {
     },
   });
 
-  registerCleanupFunction(async function () {
-    await experimentCleanup();
-  });
-
   await BrowserTestUtils.withNewTab(
     {
       gBrowser,
@@ -111,4 +107,6 @@ add_task(async function test_experiment_iconOnly() {
       });
     }
   );
+
+  experimentCleanup();
 });

toolkit/components/search/tests/xpcshell/test_search_config_v2_nimbus.js

@@ -69,7 +69,7 @@ add_task(async function test_nimbus_experiment_enabled() {
     "All application provided engines for search-config-v2 should be instances of AppProvidedSearchEngine."
   );
 
-  await doExperimentCleanup();
+  doExperimentCleanup();
 
   Assert.equal(
     Services.prefs.getBoolPref("browser.search.newSearchConfig.enabled"),