Bug 1686315 - Fix up some tests to pass with content prompt subdialogs enabled or disabled. r=Gijs

Batch 1 of 2.

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

browser/base/content/test/general/browser_beforeunload_duplicate_dialogs.js

@@ -1,10 +1,20 @@
 const TEST_PAGE =
   "http://mochi.test:8888/browser/browser/base/content/test/general/file_double_close_tab.html";
 
+const CONTENT_PROMPT_SUBDIALOG = Services.prefs.getBoolPref(
+  "prompts.contentPromptSubDialog",
+  false
+);
+
 var expectingDialog = false;
 var wantToClose = true;
 var resolveDialogPromise;
+
 function onTabModalDialogLoaded(node) {
+  ok(
+    !CONTENT_PROMPT_SUBDIALOG,
+    "Should not be using content prompt subdialogs."
+  );
   ok(expectingDialog, "Should be expecting this dialog.");
   expectingDialog = false;
   if (wantToClose) {
@@ -19,19 +29,34 @@ function onTabModalDialogLoaded(node) {
   }
 }
 
+function onCommonDialogLoaded(promptWindow) {
+  ok(CONTENT_PROMPT_SUBDIALOG, "Should be using content prompt subdialogs.");
+  ok(expectingDialog, "Should be expecting this dialog.");
+  expectingDialog = false;
+  let dialog = promptWindow.Dialog;
+  if (wantToClose) {
+    // This accepts the dialog, closing it
+    dialog.ui.button0.click();
+  } else {
+    // This keeps the page open
+    dialog.ui.button1.click();
+  }
+  if (resolveDialogPromise) {
+    resolveDialogPromise();
+  }
+}
+
 SpecialPowers.pushPrefEnv({
   set: [["dom.require_user_interaction_for_beforeunload", false]],
 });
 
-SpecialPowers.pushPrefEnv({
-  set: [["prompts.contentPromptSubDialog", false]],
-});
-
 // Listen for the dialog being created
 Services.obs.addObserver(onTabModalDialogLoaded, "tabmodal-dialog-loaded");
+Services.obs.addObserver(onCommonDialogLoaded, "common-dialog-loaded");
 registerCleanupFunction(() => {
   Services.prefs.clearUserPref("browser.tabs.warnOnClose");
   Services.obs.removeObserver(onTabModalDialogLoaded, "tabmodal-dialog-loaded");
+  Services.obs.removeObserver(onCommonDialogLoaded, "common-dialog-loaded");
 });
 
 add_task(async function closeLastTabInWindow() {

browser/base/content/test/general/browser_double_close_tab.js

@@ -4,18 +4,25 @@ const TEST_PAGE =
   "http://mochi.test:8888/browser/browser/base/content/test/general/file_double_close_tab.html";
 var testTab;
 
+const CONTENT_PROMPT_SUBDIALOG = Services.prefs.getBoolPref(
+  "prompts.contentPromptSubDialog",
+  false
+);
+
 function waitForDialog(callback) {
-  function onTabModalDialogLoaded(node) {
-    Services.obs.removeObserver(
-      onTabModalDialogLoaded,
-      "tabmodal-dialog-loaded"
-    );
+  function onDialogLoaded(nodeOrDialogWindow) {
+    let node = CONTENT_PROMPT_SUBDIALOG
+      ? nodeOrDialogWindow.document.querySelector("dialog")
+      : nodeOrDialogWindow;
+    Services.obs.removeObserver(onDialogLoaded, "tabmodal-dialog-loaded");
+    Services.obs.removeObserver(onDialogLoaded, "common-dialog-loaded");
     // Allow dialog's onLoad call to run to completion
     Promise.resolve().then(() => callback(node));
   }
 
   // Listen for the dialog being created
-  Services.obs.addObserver(onTabModalDialogLoaded, "tabmodal-dialog-loaded");
+  Services.obs.addObserver(onDialogLoaded, "tabmodal-dialog-loaded");
+  Services.obs.addObserver(onDialogLoaded, "common-dialog-loaded");
 }
 
 function waitForDialogDestroyed(node, callback) {
@@ -27,6 +34,11 @@ function waitForDialogDestroyed(node, callback) {
     }
   });
   observer.observe(node.parentNode, { childList: true });
+
+  if (CONTENT_PROMPT_SUBDIALOG) {
+    node.ownerGlobal.addEventListener("unload", done);
+  }
+
   let failureTimeout = setTimeout(function() {
     ok(false, "Dialog should have been destroyed");
     done();
@@ -36,7 +48,13 @@ function waitForDialogDestroyed(node, callback) {
     clearTimeout(failureTimeout);
     observer.disconnect();
     observer = null;
-    callback();
+
+    if (CONTENT_PROMPT_SUBDIALOG) {
+      node.ownerGlobal.removeEventListener("unload", done);
+      SimpleTest.executeSoon(callback);
+    } else {
+      callback();
+    }
   }
 }
 
@@ -45,11 +63,8 @@ add_task(async function() {
     set: [["dom.require_user_interaction_for_beforeunload", false]],
   });
 
-  await SpecialPowers.pushPrefEnv({
-    set: [["prompts.contentPromptSubDialog", false]],
-  });
-
   testTab = await BrowserTestUtils.openNewForegroundTab(gBrowser, TEST_PAGE);
+
   // XXXgijs the reason this has nesting and callbacks rather than promises is
   // that DOM promises resolve on the next tick. So they're scheduled
   // in an event queue. So when we spin a new event queue for a modal dialog...
@@ -60,14 +75,26 @@ add_task(async function() {
       waitForDialogDestroyed(dialogNode, () => {
         let doCompletion = () => setTimeout(resolveOuter, 0);
         info("Now checking if dialog is destroyed");
-        ok(!dialogNode.parentNode, "onbeforeunload dialog should be gone.");
-        if (dialogNode.parentNode) {
-          // Failed to remove onbeforeunload dialog, so do it ourselves:
-          let leaveBtn = dialogNode.querySelector(".tabmodalprompt-button0");
-          waitForDialogDestroyed(dialogNode, doCompletion);
-          EventUtils.synthesizeMouseAtCenter(leaveBtn, {});
-          return;
+
+        if (CONTENT_PROMPT_SUBDIALOG) {
+          ok(
+            !dialogNode.ownerGlobal || dialogNode.ownerGlobal.closed,
+            "onbeforeunload dialog should be gone."
+          );
+          if (dialogNode.ownerGlobal && !dialogNode.ownerGlobal.closed) {
+            dialogNode.acceptDialog();
+          }
+        } else {
+          ok(!dialogNode.parentNode, "onbeforeunload dialog should be gone.");
+          if (dialogNode.parentNode) {
+            // Failed to remove onbeforeunload dialog, so do it ourselves:
+            let leaveBtn = dialogNode.querySelector(".tabmodalprompt-button0");
+            waitForDialogDestroyed(dialogNode, doCompletion);
+            EventUtils.synthesizeMouseAtCenter(leaveBtn, {});
+            return;
+          }
         }
+
         doCompletion();
       });
       // Click again:

browser/base/content/test/pageActions/browser_page_action_menu_add_search_engine.js

@@ -107,10 +107,6 @@ add_task(async function one() {
 
 // Checks the panel button with a page that offers an invalid engine.
 add_task(async function invalid() {
-  await SpecialPowers.pushPrefEnv({
-    set: [["prompts.contentPromptSubDialog", false]],
-  });
-
   let url =
     getRootDirectory(gTestPath) +
     "page_action_menu_add_search_engine_invalid.html";

browser/base/content/test/tabPrompts/browser_beforeunload_urlbar.js

@@ -8,13 +8,15 @@ const TEST_ROOT = getRootDirectory(gTestPath).replace(
   "http://example.com"
 );
 
+const CONTENT_PROMPT_SUBDIALOG = Services.prefs.getBoolPref(
+  "prompts.contentPromptSubDialog",
+  false
+);
+
 add_task(async function test_beforeunload_stay_clears_urlbar() {
   await SpecialPowers.pushPrefEnv({
     set: [["dom.require_user_interaction_for_beforeunload", false]],
   });
-  await SpecialPowers.pushPrefEnv({
-    set: [["prompts.contentPromptSubDialog", false]],
-  });
 
   const TEST_URL = TEST_ROOT + "file_beforeunload_stop.html";
   await BrowserTestUtils.withNewTab(TEST_URL, async function(browser) {
@@ -23,18 +25,29 @@ add_task(async function test_beforeunload_stay_clears_urlbar() {
     gURLBar.inputField.value = inputValue.slice(0, -1);
     EventUtils.sendString(inputValue.slice(-1));
 
-    let promptOpenedPromise = TestUtils.topicObserved("tabmodal-dialog-loaded");
-    EventUtils.synthesizeKey("VK_RETURN");
-    await promptOpenedPromise;
-    let promptElement = browser.parentNode.querySelector("tabmodalprompt");
+    if (CONTENT_PROMPT_SUBDIALOG) {
+      let promptOpenedPromise = BrowserTestUtils.promiseAlertDialogOpen(
+        "cancel"
+      );
+      EventUtils.synthesizeKey("VK_RETURN");
+      await promptOpenedPromise;
+      await TestUtils.waitForTick();
+    } else {
+      let promptOpenedPromise = TestUtils.topicObserved(
+        "tabmodal-dialog-loaded"
+      );
+      EventUtils.synthesizeKey("VK_RETURN");
+      await promptOpenedPromise;
+      let promptElement = browser.parentNode.querySelector("tabmodalprompt");
 
-    // Click the cancel button
-    promptElement.querySelector(".tabmodalprompt-button1").click();
+      // Click the cancel button
+      promptElement.querySelector(".tabmodalprompt-button1").click();
+      await TestUtils.waitForCondition(
+        () => promptElement.parentNode == null,
+        "tabprompt should be removed"
+      );
+    }
 
-    await TestUtils.waitForCondition(
-      () => promptElement.parentNode == null,
-      "tabprompt should be removed"
-    );
     // Can't just compare directly with TEST_URL because the URL may be trimmed.
     // Just need it to not be the example.org thing we typed in.
     ok(

browser/base/content/test/tabPrompts/browser_multiplePrompts.js

@@ -1,56 +1,80 @@
 "use strict";
 
-/*
- * This test triggers multiple alerts on one single tab, because it"s possible
- * for web content to do so. The behavior is described in bug 1266353.
+const CONTENT_PROMPT_SUBDIALOG = Services.prefs.getBoolPref(
+  "prompts.contentPromptSubDialog",
+  false
+);
+
+/**
+ * Goes through a stacked series of dialogs opened with
+ * CONTENT_PROMPT_SUBDIALOG set to true, and ensures that
+ * the oldest one is front-most and has the right type. It
+ * then closes the oldest to newest dialog.
  *
- * We assert the presentation of the multiple alerts, ensuring we show only
- * the oldest one.
+ * @param {Element} tab The <tab> that has had content dialogs opened
+ * for it.
+ * @param {Number} promptCount How many dialogs we expected to have been
+ * opened.
+ *
+ * @return {Promise}
+ * @resolves {undefined} Once the dialogs have all been closed.
  */
-add_task(async function() {
-  await SpecialPowers.pushPrefEnv({
-    set: [["prompts.contentPromptSubDialog", false]],
-  });
+async function closeDialogs(tab, dialogCount) {
+  let dialogElementsCount = dialogCount;
+  let dialogs = tab.linkedBrowser.tabDialogBox.getContentDialogManager()
+    .dialogs;
 
-  const PROMPTCOUNT = 9;
+  is(
+    dialogs.length,
+    dialogElementsCount,
+    "There should be " + dialogElementsCount + " dialog(s)."
+  );
 
-  let contentScript = function(MAX_PROMPT) {
-    var i = MAX_PROMPT;
-    let fns = ["alert", "prompt", "confirm"];
-    function openDialog() {
-      i--;
-      if (i) {
-        SpecialPowers.Services.tm.dispatchToMainThread(openDialog);
-      }
-      window[fns[i % 3]](fns[i % 3] + " countdown #" + i);
-    }
-    SpecialPowers.Services.tm.dispatchToMainThread(openDialog);
-  };
-  let url =
-    "data:text/html,<script>(" +
-    encodeURIComponent(contentScript.toSource()) +
-    ")(" +
-    PROMPTCOUNT +
-    ");</script>";
+  let i = dialogElementsCount - 1;
+  for (let dialog of dialogs) {
+    dialog.focus(true);
+    await dialog._dialogReady;
 
-  let promptsOpenedPromise = new Promise(function(resolve) {
-    let unopenedPromptCount = PROMPTCOUNT;
-    Services.obs.addObserver(function observer() {
-      unopenedPromptCount--;
-      if (!unopenedPromptCount) {
-        Services.obs.removeObserver(observer, "tabmodal-dialog-loaded");
-        info("Prompts opened.");
-        resolve();
-      }
-    }, "tabmodal-dialog-loaded");
-  });
+    let dialogWindow = dialog.frameContentWindow;
+    let expectedType = ["alert", "prompt", "confirm"][i % 3];
 
-  let tab = await BrowserTestUtils.openNewForegroundTab(gBrowser, url, true);
-  info("Tab loaded");
+    is(
+      dialogWindow.Dialog.args.text,
+      expectedType + " countdown #" + i,
+      "The #" + i + " alert should be labelled as such."
+    );
+    i--;
 
-  await promptsOpenedPromise;
+    dialogWindow.Dialog.ui.button0.click();
 
-  let promptElementsCount = PROMPTCOUNT;
+    // The click is handled async; wait for an event loop turn for that to
+    // happen.
+    await new Promise(function(resolve) {
+      Services.tm.dispatchToMainThread(resolve);
+    });
+  }
+
+  dialogs = tab.linkedBrowser.tabDialogBox.getContentDialogManager().dialogs;
+  is(dialogs.length, 0, "Dialogs should all be dismissed.");
+}
+
+/**
+ * Goes through a stacked series of tabprompt modals opened with
+ * CONTENT_PROMPT_SUBDIALOG set to false, and ensures that
+ * the oldest one is front-most and has the right type. It also
+ * ensures that the other tabprompt modals are hidden. It
+ * then closes the oldest to newest dialog.
+ *
+ * @param {Element} tab The <tab> that has had tabprompt modals opened
+ * for it.
+ * @param {Number} promptCount How many modals we expected to have been
+ * opened.
+ *
+ * @return {Promise}
+ * @resolves {undefined} Once the modals have all been closed.
+ */
+async function closeTabModals(tab, promptCount) {
+  let promptElementsCount = promptCount;
   while (promptElementsCount--) {
     let promptElements = tab.linkedBrowser.parentNode.querySelectorAll(
       "tabmodalprompt"
@@ -62,6 +86,7 @@ add_task(async function() {
     );
     // The oldest should be the first.
     let i = 0;
+
     for (let promptElement of promptElements) {
       let prompt = tab.linkedBrowser.tabModalPromptBox.getPrompt(promptElement);
       let expectedType = ["alert", "prompt", "confirm"][i % 3];
@@ -91,6 +116,57 @@ add_task(async function() {
     "tabmodalprompt"
   );
   is(promptElements.length, 0, "Prompts should all be dismissed.");
+}
+
+/*
+ * This test triggers multiple alerts on one single tab, because it"s possible
+ * for web content to do so. The behavior is described in bug 1266353.
+ *
+ * We assert the presentation of the multiple alerts, ensuring we show only
+ * the oldest one.
+ */
+add_task(async function() {
+  const PROMPTCOUNT = 9;
+
+  let unopenedPromptCount = PROMPTCOUNT;
+
+  let tab = await BrowserTestUtils.openNewForegroundTab(
+    gBrowser,
+    "http://example.com",
+    true
+  );
+  info("Tab loaded");
+
+  let promptsOpenedPromise = BrowserTestUtils.waitForEvent(
+    tab.linkedBrowser,
+    "DOMWillOpenModalDialog",
+    false,
+    () => {
+      unopenedPromptCount--;
+      return unopenedPromptCount == 0;
+    }
+  );
+
+  await SpecialPowers.spawn(tab.linkedBrowser, [PROMPTCOUNT], maxPrompts => {
+    var i = maxPrompts;
+    let fns = ["alert", "prompt", "confirm"];
+    function openDialog() {
+      i--;
+      if (i) {
+        SpecialPowers.Services.tm.dispatchToMainThread(openDialog);
+      }
+      content[fns[i % 3]](fns[i % 3] + " countdown #" + i);
+    }
+    SpecialPowers.Services.tm.dispatchToMainThread(openDialog);
+  });
+
+  await promptsOpenedPromise;
+
+  if (CONTENT_PROMPT_SUBDIALOG) {
+    await closeDialogs(tab, PROMPTCOUNT);
+  } else {
+    await closeTabModals(tab, PROMPTCOUNT);
+  }
 
   BrowserTestUtils.removeTab(tab);
 });

browser/base/content/test/tabPrompts/browser_openPromptInBackgroundTab.js

@@ -22,6 +22,7 @@ registerCleanupFunction(function() {
  * checking the checkbox does actually enable that behaviour.
  */
 add_task(async function test_old_modal_ui() {
+  // We're intentionally testing the old modal mechanism, so disable the new one.
   await SpecialPowers.pushPrefEnv({
     set: [["prompts.contentPromptSubDialog", false]],
   });
@@ -130,6 +131,7 @@ add_task(async function test_old_modal_ui() {
 });
 
 add_task(async function test_new_modal_ui() {
+  // We're intentionally testing the new modal mechanism, so make sure it's enabled.
   await SpecialPowers.pushPrefEnv({
     set: [["prompts.contentPromptSubDialog", true]],
   });

browser/base/content/test/tabs/browser_close_during_beforeunload.js

@@ -4,17 +4,20 @@
 // beforeunload confirmation ignores the beforeunload listener and
 // unblocks the original close call.
 
-const DIALOG_TOPIC = "tabmodal-dialog-loaded";
+const CONTENT_PROMPT_SUBDIALOG = Services.prefs.getBoolPref(
+  "prompts.contentPromptSubDialog",
+  false
+);
+
+const DIALOG_TOPIC = CONTENT_PROMPT_SUBDIALOG
+  ? "common-dialog-loaded"
+  : "tabmodal-dialog-loaded";
 
 add_task(async function() {
   await SpecialPowers.pushPrefEnv({
     set: [["dom.require_user_interaction_for_beforeunload", false]],
   });
 
-  await SpecialPowers.pushPrefEnv({
-    set: [["prompts.contentPromptSubDialog", false]],
-  });
-
   let win = await BrowserTestUtils.openNewBrowserWindow();
 
   let browser = win.gBrowser.selectedBrowser;

browser/components/extensions/test/browser/browser_ext_tabs_attention.js

@@ -3,10 +3,6 @@
 "use strict";
 
 add_task(async function tabsAttention() {
-  await SpecialPowers.pushPrefEnv({
-    set: [["prompts.contentPromptSubDialog", false]],
-  });
-
   let tab1 = await BrowserTestUtils.openNewForegroundTab(
     gBrowser,
     "http://example.com/?2",

browser/components/search/test/browser/browser_searchbar_addEngine.js

@@ -15,10 +15,6 @@ const searchPopup = document.getElementById("PopupSearchAutoComplete");
 let searchbar;
 
 add_task(async function setup() {
-  await SpecialPowers.pushPrefEnv({
-    set: [["prompts.contentPromptSubDialog", false]],
-  });
-
   searchbar = await gCUITestUtils.addSearchBar();
 
   registerCleanupFunction(async function() {

devtools/server/tests/browser/browser_navigateEvents.js

@@ -6,16 +6,16 @@
 const URL1 = MAIN_DOMAIN + "navigate-first.html";
 const URL2 = MAIN_DOMAIN + "navigate-second.html";
 
+const { PromptTestUtils } = ChromeUtils.import(
+  "resource://testing-common/PromptTestUtils.jsm"
+);
+
 var isE10s = Services.appinfo.browserTabsRemoteAutostart;
 
 SpecialPowers.pushPrefEnv({
   set: [["dom.require_user_interaction_for_beforeunload", false]],
 });
 
-SpecialPowers.pushPrefEnv({
-  set: [["prompts.contentPromptSubDialog", false]],
-});
-
 var signalAllEventsReceived;
 var onAllEventsReceived = new Promise(resolve => {
   signalAllEventsReceived = resolve;
@@ -87,22 +87,6 @@ function assertEvent(event, data) {
   }
 }
 
-function waitForOnBeforeUnloadDialog(browser, callback) {
-  browser.addEventListener(
-    "DOMWillOpenModalDialog",
-    async function(event) {
-      const stack = browser.parentNode;
-      const dialogs = stack.getElementsByTagName("tabmodalprompt");
-      await waitUntil(() => dialogs[0]);
-      const { button0, button1 } = browser.tabModalPromptBox.getPrompt(
-        dialogs[0]
-      ).ui;
-      callback(button0, button1);
-    },
-    { capture: true, once: true }
-  );
-}
-
 var httpObserver = function(subject, topic, state) {
   const channel = subject.QueryInterface(Ci.nsIHttpChannel);
   const url = channel.URI.spec;
@@ -141,10 +125,12 @@ add_task(async function() {
   const browser = await addTab(URL1);
 
   // Listen for alert() call being made in navigate-first during unload
-  waitForOnBeforeUnloadDialog(browser, function(btnLeave, btnStay) {
+  const beforeUnloadPromise = PromptTestUtils.handleNextPrompt(
+    browser,
+    { modalType: Services.prompt.MODAL_TYPE_CONTENT, promptType: "confirmEx" },
+    { buttonNumClick: 0 }
+  ).then(() => {
     assertEvent("unload-dialog");
-    // accept to quit this page to another
-    btnLeave.click();
   });
 
   // Listen for messages sent by the content task
@@ -184,6 +170,7 @@ add_task(async function() {
     gBrowser.selectedBrowser
   );
   BrowserTestUtils.loadURI(gBrowser.selectedBrowser, URL2);
+  await beforeUnloadPromise;
   await onBrowserLoaded;
 
   // Wait for all events to be received

docshell/test/browser/browser_bug1415918_beforeunload_options.js

@@ -6,6 +6,10 @@ const TEST_PATH = getRootDirectory(gTestPath).replace(
   "http://example.com"
 );
 
+const { PromptTestUtils } = ChromeUtils.import(
+  "resource://testing-common/PromptTestUtils.jsm"
+);
+
 SimpleTest.requestFlakyTimeout("Needs to test a timeout");
 
 function delay(msec) {
@@ -13,15 +17,27 @@ function delay(msec) {
   return new Promise(resolve => setTimeout(resolve, msec));
 }
 
+function allowNextNavigation(browser) {
+  return PromptTestUtils.handleNextPrompt(
+    browser,
+    { modalType: Services.prompt.MODAL_TYPE_CONTENT, promptType: "confirmEx" },
+    { buttonNumClick: 0 }
+  );
+}
+
+function cancelNextNavigation(browser) {
+  return PromptTestUtils.handleNextPrompt(
+    browser,
+    { modalType: Services.prompt.MODAL_TYPE_CONTENT, promptType: "confirmEx" },
+    { buttonNumClick: 1 }
+  );
+}
+
 add_task(async function test() {
   await SpecialPowers.pushPrefEnv({
     set: [["dom.require_user_interaction_for_beforeunload", false]],
   });
 
-  await SpecialPowers.pushPrefEnv({
-    set: [["prompts.contentPromptSubDialog", false]],
-  });
-
   const permitUnloadTimeout = Services.prefs.getIntPref(
     "dom.beforeunload_timeout_ms"
   );
@@ -36,47 +52,28 @@ add_task(async function test() {
     });
   });
 
-  let allowNavigation;
-  let promptShown = false;
-  let promptDismissed = false;
-  let promptTimeout;
-
-  const DIALOG_TOPIC = "tabmodal-dialog-loaded";
-  async function observer(node) {
-    promptShown = true;
-
-    if (promptTimeout) {
-      await delay(promptTimeout);
-    }
-
-    let button = node.querySelector(
-      `.tabmodalprompt-button${allowNavigation ? 0 : 1}`
-    );
-    button.click();
-    promptDismissed = true;
-  }
-  Services.obs.addObserver(observer, DIALOG_TOPIC);
-
   /*
    * Check condition where beforeunload handlers request a prompt.
    */
 
   // Prompt is shown, user clicks OK.
-  allowNavigation = true;
-  promptShown = false;
 
+  let promptShownPromise = allowNextNavigation(browser);
   ok(browser.permitUnload().permitUnload, "permit unload should be true");
-  ok(promptShown, "prompt should have been displayed");
+  await promptShownPromise;
 
   // Prompt is shown, user clicks CANCEL.
-  allowNavigation = false;
-  promptShown = false;
-
+  promptShownPromise = cancelNextNavigation(browser);
   ok(!browser.permitUnload().permitUnload, "permit unload should be false");
-  ok(promptShown, "prompt should have been displayed");
+  await promptShownPromise;
 
   // Prompt is not shown, don't permit unload.
-  promptShown = false;
+  let promptShown = false;
+  let shownCallback = () => {
+    promptShown = true;
+  };
+
+  browser.addEventListener("DOMWillOpenModalDialog", shownCallback);
   ok(
     !browser.permitUnload("dontUnload").permitUnload,
     "permit unload should be false"
@@ -90,10 +87,20 @@ add_task(async function test() {
     "permit unload should be true"
   );
   ok(!promptShown, "prompt should not have been displayed");
+  browser.removeEventListener("DOMWillOpenModalDialog", shownCallback);
+
+  promptShownPromise = PromptTestUtils.waitForPrompt(browser, {
+    modalType: Services.prompt.MODAL_TYPE_CONTENT,
+    promptType: "confirmEx",
+  });
+
+  let promptDismissed = false;
+  let closedCallback = () => {
+    promptDismissed = true;
+  };
+
+  browser.addEventListener("DOMModalDialogClosed", closedCallback);
 
-  promptShown = false;
-  promptDismissed = false;
-  promptTimeout = 3 * permitUnloadTimeout;
   let promise = browser.asyncPermitUnload();
 
   let promiseResolved = false;
@@ -101,8 +108,7 @@ add_task(async function test() {
     promiseResolved = true;
   });
 
-  await TestUtils.waitForCondition(() => promptShown);
-  ok(!promptDismissed, "Should not have dismissed prompt yet");
+  let dialog = await promptShownPromise;
   ok(!promiseResolved, "Should not have resolved promise yet");
 
   await delay(permitUnloadTimeout * 1.5);
@@ -110,24 +116,28 @@ add_task(async function test() {
   ok(!promptDismissed, "Should not have dismissed prompt yet");
   ok(!promiseResolved, "Should not have resolved promise yet");
 
+  await PromptTestUtils.handlePrompt(dialog, { buttonNumClick: 1 });
+
   let { permitUnload } = await promise;
   ok(promptDismissed, "Should have dismissed prompt");
   ok(!permitUnload, "Should not have permitted unload");
 
-  promptTimeout = null;
+  browser.removeEventListener("DOMModalDialogClosed", closedCallback);
+
+  promptShownPromise = allowNextNavigation(browser);
 
   /*
    * Check condition where no one requests a prompt.  In all cases,
    * permitUnload should be true, and all handlers fired.
    */
-
-  allowNavigation = true;
-
   url += "?1";
   BrowserTestUtils.loadURI(browser, url);
   await BrowserTestUtils.browserLoaded(browser, false, url);
+  await promptShownPromise;
 
   promptShown = false;
+  browser.addEventListener("DOMWillOpenModalDialog", shownCallback);
+
   ok(browser.permitUnload().permitUnload, "permit unload should be true");
   ok(!promptShown, "prompt should not have been displayed");
 
@@ -145,7 +155,7 @@ add_task(async function test() {
   );
   ok(!promptShown, "prompt should not have been displayed");
 
-  await BrowserTestUtils.removeTab(tab);
+  browser.removeEventListener("DOMWillOpenModalDialog", shownCallback);
 
-  Services.obs.removeObserver(observer, DIALOG_TOPIC);
+  await BrowserTestUtils.removeTab(tab, { skipPermitUnload: true });
 });