Bug 1613994 - Wait for state telemetry event after resetting enabled pref to ensure add-on has a chance to react. r=dragana

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

--HG--
extra : moz-landing-system : lando

browser/extensions/doh-rollout/background.js

@@ -73,21 +73,25 @@ const stateManager = {
       case "uninstalled":
         break;
       case "disabled":
-        rollout.setSetting(TRR_MODE_PREF, 0);
+        await rollout.setSetting(TRR_MODE_PREF, 0);
         break;
       case "manuallyDisabled":
-        browser.experiments.preferences.clearUserPref(DOH_SELF_ENABLED_PREF);
+        await browser.experiments.preferences.clearUserPref(
+          DOH_SELF_ENABLED_PREF
+        );
         break;
       case "UIOk":
-        rollout.setSetting(DOH_SELF_ENABLED_PREF, true);
+        await rollout.setSetting(DOH_SELF_ENABLED_PREF, true);
         break;
       case "enabled":
-        rollout.setSetting(TRR_MODE_PREF, 2);
-        rollout.setSetting(DOH_SELF_ENABLED_PREF, true);
+        await rollout.setSetting(TRR_MODE_PREF, 2);
+        await rollout.setSetting(DOH_SELF_ENABLED_PREF, true);
         break;
       case "UIDisabled":
-        rollout.setSetting(TRR_MODE_PREF, 5);
-        browser.experiments.preferences.clearUserPref(DOH_SELF_ENABLED_PREF);
+        await rollout.setSetting(TRR_MODE_PREF, 5);
+        await browser.experiments.preferences.clearUserPref(
+          DOH_SELF_ENABLED_PREF
+        );
         break;
     }
 

browser/extensions/doh-rollout/test/browser/browser_rollback.js

@@ -67,6 +67,7 @@ add_task(async function testRollback() {
   // Rollback!
   setPassingHeuristics();
   Preferences.reset(prefs.DOH_ENABLED_PREF);
+  await waitForStateTelemetry();
   await ensureTRRMode(0);
   await ensureNoHeuristicsTelemetry();
   simulateNetworkChange();
@@ -87,6 +88,7 @@ add_task(async function testRollback() {
 
   // Rollback again for good measure! This time with failing heuristics.
   Preferences.reset(prefs.DOH_ENABLED_PREF);
+  await waitForStateTelemetry();
   await ensureNoTRRModeChange(0);
   await ensureNoHeuristicsTelemetry();
   simulateNetworkChange();
@@ -107,6 +109,7 @@ add_task(async function testRollback() {
 
   // Rollback again, this time with TRR mode set to 2 prior to doing so.
   Preferences.reset(prefs.DOH_ENABLED_PREF);
+  await waitForStateTelemetry();
   await ensureTRRMode(0);
   await ensureNoHeuristicsTelemetry();
   simulateNetworkChange();

browser/extensions/doh-rollout/test/browser/head.js

@@ -117,6 +117,19 @@ function ensureNoHeuristicsTelemetry() {
   is(events.length, 0, "Found no heuristics events.");
 }
 
+async function waitForStateTelemetry() {
+  let events;
+  await BrowserTestUtils.waitForCondition(() => {
+    events = Services.telemetry.snapshotEvents(
+      Ci.nsITelemetry.DATASET_PRERELEASE_CHANNELS
+    ).dynamic;
+    return events;
+  });
+  events = events.filter(e => e[1] == "doh" && e[2] == "state");
+  is(events.length, 1, "Found the expected state event.");
+  Services.telemetry.clearEvents();
+}
+
 function setPassingHeuristics() {
   Preferences.set(prefs.MOCK_HEURISTICS_PREF, fakePassingHeuristics);
 }