Bug 1736103 - fix download file extension correction test some more for the no-download-modal case, r=sfoster

On try and on my machine, after D128648 the test still times out because
for the test_broken_saved_handlerinfo_and_useless_mimetypes test, first
awaiting the `task` that gets passed means that the download has already
finished by the time we create the download finished promise.

This patch moves creating that promise before the `task` is started.
It also broadens the conditions in which we await this promise, because
in the case without file extensions and without a dialog, we were otherwise
no longer awaiting anything, which meant that the next test's check for
a finished download would resolve with the previous test's download, causing
confusion/breakage.

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

uriloader/exthandler/tests/mochitest/browser_extension_correction.js

@@ -44,7 +44,7 @@ async function checkDownloadWithExtensionState(
   task,
   { type, shouldHaveExtension, expectedName = null }
 ) {
-  const shouldExpectDialog = !SpecialPowers.Services.prefs.getBoolPref(
+  const shouldExpectDialog = !Services.prefs.getBoolPref(
     "browser.download.improvements_to_download_panel",
     false
   );
@@ -54,6 +54,13 @@ async function checkDownloadWithExtensionState(
     winPromise = BrowserTestUtils.domWindowOpenedAndLoaded();
   }
 
+  let publicList = await Downloads.getList(Downloads.PUBLIC);
+  let shouldCheckFilename = shouldHaveExtension || !shouldExpectDialog;
+
+  let downloadFinishedPromise = shouldCheckFilename
+    ? promiseDownloadFinished(publicList)
+    : null;
+
   await task();
 
   let win;
@@ -78,22 +85,19 @@ async function checkDownloadWithExtensionState(
     }
   }
 
-  if (shouldHaveExtension) {
-    // Wait for the download.
-    let publicList = await Downloads.getList(Downloads.PUBLIC);
-    let downloadFinishedPromise = promiseDownloadFinished(publicList);
-
+  if (shouldExpectDialog && shouldHaveExtension) {
     // Then pick "save" in the dialog, if we have a dialog.
-    if (shouldExpectDialog) {
-      let dialog = win.document.getElementById("unknownContentType");
-      win.document.getElementById("save").click();
-      let button = dialog.getButton("accept");
-      button.disabled = false;
-      dialog.acceptDialog();
-    }
+    let dialog = win.document.getElementById("unknownContentType");
+    win.document.getElementById("save").click();
+    let button = dialog.getButton("accept");
+    button.disabled = false;
+    dialog.acceptDialog();
+  }
 
-    // Wait for the download to finish and check the extension is correct.
-    let download = await downloadFinishedPromise;
+  // Wait for the download if it exists (may produce null).
+  let download = await downloadFinishedPromise;
+  if (download) {
+    // Check the download's extension is correct.
     is(
       PathUtils.filename(download.target.path),
       expectedName,