Bug 1781397 - Fix the intermittent issue of test_bug1254766. r=dimi

To fix the intermittent issue, the patch ensures the safebrowsing
database is initiated after reload the database before we load the
test frame.

It also fixes the issue that classifierHelper doesn't properly clean up
the prefs used in the test.

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

toolkit/components/url-classifier/tests/mochitest/classifierCommon.js

@@ -77,6 +77,20 @@ function waitForInit() {
   }
 }
 
+function doGetTables() {
+  const callback = tables => {
+    sendAsyncMessage("GetTableSuccess", tables);
+  };
+
+  try {
+    dbService.getTables(callback);
+  } catch (e) {
+    setTimeout(() => {
+      doGetTables();
+    }, 1000);
+  }
+}
+
 addMessageListener("doUpdate", ({ testUpdate }) => {
   doUpdate(testUpdate);
 });
@@ -88,3 +102,7 @@ addMessageListener("doReload", () => {
 addMessageListener("waitForInit", () => {
   waitForInit();
 });
+
+addMessageListener("doGetTables", () => {
+  doGetTables();
+});

toolkit/components/url-classifier/tests/mochitest/classifierHelper.js

@@ -4,6 +4,7 @@ if (typeof classifierHelper == "undefined") {
 
 const CLASSIFIER_COMMON_URL = SimpleTest.getTestFileURL("classifierCommon.js");
 var gScript = SpecialPowers.loadChromeScript(CLASSIFIER_COMMON_URL);
+var gOriginalGetHashURL;
 
 const PREFS = {
   PROVIDER_LISTS: "browser.safebrowsing.provider.mozilla.lists",
@@ -49,6 +50,9 @@ classifierHelper.allowCompletion = async function (lists, url) {
     await SpecialPowers.setCharPref(PREFS.DISALLOW_COMPLETIONS, pref);
   }
 
+  // Store the original get hash URL in order to reset it back during clean up.
+  gOriginalGetHashURL = SpecialPowers.getCharPref(PREFS.PROVIDER_GETHASHURL);
+
   // Set get hash url
   await SpecialPowers.setCharPref(PREFS.PROVIDER_GETHASHURL, url);
 };
@@ -121,6 +125,17 @@ classifierHelper.reloadDatabase = function () {
   });
 };
 
+classifierHelper.getTables = function () {
+  return new Promise(resolve => {
+    gScript.addMessageListener("GetTableSuccess", function handler(tables) {
+      gScript.removeMessageListener("GetTableSuccess", handler);
+      resolve(tables);
+    });
+
+    gScript.sendAsyncMessage("doGetTables");
+  });
+};
+
 classifierHelper._update = function (testUpdate, onsuccess, onerror) {
   // Queue the task if there is still an on-going update
   classifierHelper._updates.push({
@@ -173,9 +188,11 @@ classifierHelper._setup = function () {
 
 classifierHelper._cleanup = function () {
   // clean all the preferences may touch by helper
-  for (var pref in PREFS) {
-    SpecialPowers.clearUserPref(pref);
-  }
+  Object.values(PREFS).map(pref => SpecialPowers.clearUserPref(pref));
+
+  // Set the getHashURL back, the original value isn't the same as the default
+  // pref value.
+  SpecialPowers.setCharPref(PREFS.PROVIDER_GETHASHURL, gOriginalGetHashURL);
 
   if (!classifierHelper._updatesToCleanup) {
     return Promise.resolve();

toolkit/components/url-classifier/tests/mochitest/test_bug1254766.html

@@ -191,6 +191,9 @@ function testUpdateCompletionsAfterReload() {
   return Promise.resolve()
     .then(addCompletionToDB)
     .then(classifierHelper.reloadDatabase)
+    // Call getTables to ensure the DB is fully reloaded before we load the test
+    // frame.
+    .then(classifierHelper.getTables)
     .then(loadTestFrame)
     .then(() => {
       ok(gCurGethashCounter == gPreGethashCounter, "Gethash request is not triggered.");
@@ -217,6 +220,9 @@ function testGethashCompletionsAfterReload() {
       ok(gCurGethashCounter == gPreGethashCounter, "Gethash request is not triggered.");
 })
     .then(classifierHelper.reloadDatabase)
+    // Call getTables to ensure the DB is fully reloaded before we load the test
+    // frame.
+    .then(classifierHelper.getTables)
     .then(loadTestFrame)
     .then(() => {
        ok(gCurGethashCounter > gPreGethashCounter, "Gethash request is triggered.");