Bug 1896002 - Enable screenshots preview buttons if download fails. r=sfoster,kcochrane

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

browser/components/screenshots/ScreenshotsUtils.sys.mjs

@@ -1224,14 +1224,19 @@ export var ScreenshotsUtils = {
    * @param dataUrl The image data
    * @param browser The current browser
    * @param data Telemetry data
+   * @returns true if the download succeeds, otherwise false
    */
   async downloadScreenshot(title, dataUrl, browser, data) {
     // Guard against missing image data.
     if (!dataUrl) {
-      return;
+      return false;
     }
 
-    let filename = await getFilename(title, browser);
+    let { filename, accepted } = await getFilename(title, browser);
+
+    if (!accepted) {
+      return false;
+    }
 
     const targetFile = new lazy.FileUtils.File(filename);
 
@@ -1259,6 +1264,8 @@ export var ScreenshotsUtils = {
           new Blob([filename]).size
         })`
       );
+
+      return false;
     }
 
     let extra = await this.getActor(browser).sendQuery(
@@ -1274,6 +1281,8 @@ export var ScreenshotsUtils = {
       SCREENSHOTS_LAST_SAVED_METHOD_PREF,
       "download"
     );
+
+    return true;
   },
 
   recordTelemetryEvent(type, object, args) {

browser/components/screenshots/fileHelpers.mjs

@@ -97,11 +97,11 @@ export async function getFilename(filenameTitle, browser) {
     };
     let accepted = await promiseTargetFile(fpParams, browser.ownerGlobal);
     if (!accepted) {
-      return null;
+      return { filename: null, accepted };
     }
     filename = fpParams.file.path;
   }
-  return filename;
+  return { filename, accepted: true };
 }
 
 /**

browser/components/screenshots/screenshots-preview.mjs

@@ -162,8 +162,16 @@ class ScreenshotsPreview extends MozLitElement {
   }
 
   /**
-   * Disable all the buttons. After a button click, we will always close the
-   * window so there is no need to re-enable the buttons.
+   * Enable all the buttons. This will only happen when the download button is
+   * clicked and the file picker is closed without saving the image.
+   */
+  enableButtons() {
+    this.buttons.forEach(button => (button.disabled = false));
+  }
+
+  /**
+   * Disable all the buttons so they can't be clicked multiple times before
+   * successfully copying or downloading the image.
    */
   disableButtons() {
     this.buttons.forEach(button => (button.disabled = true));
@@ -176,13 +184,18 @@ class ScreenshotsPreview extends MozLitElement {
 
     // Wait for the image to be loaded before we save it
     let imageSrc = await this.imageLoadedPromise();
-    await lazy.ScreenshotsUtils.downloadScreenshot(
+    let downloadSucceeded = await lazy.ScreenshotsUtils.downloadScreenshot(
       null,
       imageSrc,
       this.openerBrowser,
       { object: "preview_download" }
     );
-    this.close();
+
+    if (downloadSucceeded) {
+      this.close();
+    } else {
+      this.enableButtons();
+    }
   }
 
   async saveToClipboard() {

browser/components/screenshots/tests/browser/browser_screenshots_test_downloads.js

@@ -46,6 +46,18 @@ function waitForFilePicker() {
     MockFilePicker.showCallback = () => {
       MockFilePicker.showCallback = null;
       ok(true, "Saw the file picker");
+
+      resolve();
+    };
+  });
+}
+
+function waitForFilePickerCancel() {
+  return new Promise(resolve => {
+    MockFilePicker.showCallback = () => {
+      MockFilePicker.showCallback = null;
+      ok(true, "Saw the file picker");
+      MockFilePicker.returnValue = MockFilePicker.returnCancel;
       resolve();
     };
   });
@@ -181,3 +193,50 @@ add_task(async function test_download_with_filepicker() {
     }
   );
 });
+
+add_task(async function test_download_filepicker_canceled() {
+  await SpecialPowers.pushPrefEnv({
+    set: [["browser.download.useDownloadDir", false]],
+  });
+
+  await BrowserTestUtils.withNewTab(
+    {
+      gBrowser,
+      url: TEST_PAGE,
+    },
+    async browser => {
+      let helper = new ScreenshotsHelper(browser);
+
+      helper.triggerUIFromToolbar();
+      await helper.waitForOverlay();
+
+      let screenshotReady = TestUtils.topicObserved(
+        "screenshots-preview-ready"
+      );
+
+      let visiblepageButton = await helper.getPanelButton("#visible-page");
+      visiblepageButton.click();
+      await screenshotReady;
+
+      let downloadButton = helper.getDialogButton("download");
+
+      let filePickerCanceled = waitForFilePickerCancel();
+      downloadButton.click();
+      info("download button clicked");
+      await filePickerCanceled;
+
+      let cancelButton = helper.getDialogButton("cancel");
+
+      await BrowserTestUtils.waitForMutationCondition(
+        cancelButton,
+        { attributes: true },
+        () => !cancelButton.disabled
+      );
+
+      let screenshotExit = TestUtils.topicObserved("screenshots-exit");
+      cancelButton.click();
+
+      await screenshotExit;
+    }
+  );
+});