1 files changed, 27 insertions, 17 deletions
diff --git a/browser/components/search/test/browser/browser_contentSearchUI.js b/browser/components/search/test/browser/browser_contentSearchUI.js
index 9196b1355c..da6044f35f 100644
--- a/browser/components/search/test/browser/browser_contentSearchUI.js
+++ b/browser/components/search/test/browser/browser_contentSearchUI.js
@@ -24,17 +24,6 @@ ChromeUtils.defineESModuleGetters(this, {
     "resource://gre/modules/SearchSuggestionController.sys.mjs",
 });
 
-const pageURL = getRootDirectory(gTestPath) + TEST_PAGE_BASENAME;
-BrowserTestUtils.registerAboutPage(
-  registerCleanupFunction,
-  "test-about-content-search-ui",
-  pageURL,
-  Ci.nsIAboutModule.URI_SAFE_FOR_UNTRUSTED_CONTENT |
-    Ci.nsIAboutModule.URI_MUST_LOAD_IN_CHILD |
-    Ci.nsIAboutModule.ALLOW_SCRIPT |
-    Ci.nsIAboutModule.URI_CAN_LOAD_IN_PRIVILEGEDABOUT_PROCESS
-);
-
 requestLongerTimeout(2);
 
 function waitForSuggestions() {
@@ -261,6 +250,19 @@ let extension1;
 let extension2;
 
 add_setup(async function () {
+  const pageURL = getRootDirectory(gTestPath) + TEST_PAGE_BASENAME;
+
+  let cleanupAboutPage;
+  await BrowserTestUtils.registerAboutPage(
+    callback => (cleanupAboutPage = callback),
+    "test-about-content-search-ui",
+    pageURL,
+    Ci.nsIAboutModule.URI_SAFE_FOR_UNTRUSTED_CONTENT |
+      Ci.nsIAboutModule.URI_MUST_LOAD_IN_CHILD |
+      Ci.nsIAboutModule.ALLOW_SCRIPT |
+      Ci.nsIAboutModule.URI_CAN_LOAD_IN_PRIVILEGEDABOUT_PROCESS
+  );
+
   let originalOnMessageSearch = ContentSearch._onMessageSearch;
   let originalOnMessageManageEngines = ContentSearch._onMessageManageEngines;
 
@@ -290,8 +292,20 @@ add_setup(async function () {
   }
 
   registerCleanupFunction(async () => {
+    // Ensure tabs are closed before we continue on with the cleanup.
+    for (let tab of tabs) {
+      BrowserTestUtils.removeTab(tab);
+    }
+    Services.search.restoreDefaultEngines();
+
+    await TestUtils.waitForTick();
+
     ContentSearch._onMessageSearch = originalOnMessageSearch;
     ContentSearch._onMessageManageEngines = originalOnMessageManageEngines;
+
+    if (cleanupAboutPage) {
+      await cleanupAboutPage();
+    }
   });
 
   await promiseTab();
@@ -1096,10 +1110,6 @@ add_task(async function settings() {
   await msg("reset");
 });
 
-add_task(async function cleanup() {
-  Services.search.restoreDefaultEngines();
-});
-
 function checkState(
   actualState,
   expectedInputVal,
@@ -1147,10 +1157,10 @@ function checkState(
 }
 
 var gMsgMan;
-
+var tabs = [];
 async function promiseTab() {
   let tab = await BrowserTestUtils.openNewForegroundTab(gBrowser);
-  registerCleanupFunction(() => BrowserTestUtils.removeTab(tab));
+  tabs.push(tab);
 
   let loadedPromise = BrowserTestUtils.firstBrowserLoaded(window);
   openTrustedLinkIn("about:test-about-content-search-ui", "current");