Bug 1171708 - Fix SessionStore tests to account for async window flushing. r=billm

browser/components/sessionstore/SessionStore.jsm

@@ -1299,7 +1299,6 @@ var SessionStoreInternal = {
         // save the state without this window to disk
         this.saveStateDelayed();
       });
-
     } else {
       this.cleanUpWindow(aWindow, winData);
     }

browser/components/sessionstore/test/browser_248970_b_perwindowpb.js

@@ -26,10 +26,10 @@ function test() {
     "//input[@type='file']":      filePath
   };
 
-  registerCleanupFunction(function() {
+  registerCleanupFunction(function* () {
-    windowsToClose.forEach(function(win) {
+    for (let win of windowsToClose) {
-      win.close();
+      yield BrowserTestUtils.closeWindow(win);
-    });
+    }
   });
 
   function test(aLambda) {

browser/components/sessionstore/test/browser_394759_basic.js

@@ -32,10 +32,7 @@ function test() {
 
     let browser = newWin.gBrowser.selectedBrowser;
     setInputChecked(browser, {id: "chk", checked: true}).then(() => {
-      newWin.close();
+      BrowserTestUtils.closeWindow(newWin).then(() => {
-
-      // Now give it time to close
-      executeSoon(function() {
         is(ss.getClosedWindowCount(), 1,
            "The closed window was added to Recently Closed Windows");
         let data = JSON.parse(ss.getClosedWindowData())[0];
@@ -75,8 +72,7 @@ function test() {
              "The window correctly restored the data associated with it");
 
           // clean up
-          newWin2.close();
+          BrowserTestUtils.closeWindow(newWin2).then(finish);
-          finish();
         }, true);
       });
     });

browser/components/sessionstore/test/browser_394759_behavior.js

@@ -1,64 +1,65 @@
-/* This Source Code Form is subject to the terms of the Mozilla Public
+/**
- * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * Test helper function that opens a series of windows, closes them
- * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+ * and then checks the closed window data from SessionStore against
+ * expected results.
+ *
+ * @param windowsToOpen (Array)
+ *        An array of Objects, where each object must define a single
+ *        property "isPopup" for whether or not the opened window should
+ *        be a popup.
+ * @param expectedResults (Array)
+ *        An Object with two properies: mac and other, where each points
+ *        at yet another Object, with the following properties:
+ *
+ *        popup (int):
+ *          The number of popup windows we expect to be in the closed window
+ *          data.
+ *        normal (int):
+ *          The number of normal windows we expect to be in the closed window
+ *          data.
+ * @returns Promise
+ */
+function testWindows(windowsToOpen, expectedResults) {
+  return Task.spawn(function*() {
+    for (let winData of windowsToOpen) {
+      let features = "chrome,dialog=no," +
+                     (winData.isPopup ? "all=no" : "all");
+      let url = "http://example.com/?window=" + windowsToOpen.length;
 
-function test() {
+      let openWindowPromise = BrowserTestUtils.waitForNewWindow();
+      openDialog(getBrowserURL(), "", features, url);
+      let win = yield openWindowPromise;
+      yield BrowserTestUtils.browserLoaded(win.gBrowser.selectedBrowser);
+
+      if (win.gMultiProcessBrowser) {
+        let tab = win.gBrowser.selectedTab;
+        yield promiseTabRestored(tab);
+      }
+
+      yield BrowserTestUtils.closeWindow(win);
+    }
+
+    let closedWindowData = JSON.parse(ss.getClosedWindowData());
+    let numPopups = closedWindowData.filter(function(el, i, arr) {
+      return el.isPopup;
+    }).length;
+    let numNormal = ss.getClosedWindowCount() - numPopups;
+    // #ifdef doesn't work in browser-chrome tests, so do a simple regex on platform
+    let oResults = navigator.platform.match(/Mac/) ? expectedResults.mac
+                                                   : expectedResults.other;
+    is(numPopups, oResults.popup,
+       "There were " + oResults.popup + " popup windows to reopen");
+    is(numNormal, oResults.normal,
+       "There were " + oResults.normal + " normal windows to repoen");
+  });
+}
+
+add_task(function* test_closed_window_states() {
   // This test takes quite some time, and timeouts frequently, so we require
   // more time to run.
   // See Bug 518970.
   requestLongerTimeout(2);
 
-  waitForExplicitFinish();
-
-  // helper function that does the actual testing
-  function openWindowRec(windowsToOpen, expectedResults, recCallback) {
-    // do actual checking
-    if (!windowsToOpen.length) {
-      let closedWindowData = JSON.parse(ss.getClosedWindowData());
-      let numPopups = closedWindowData.filter(function(el, i, arr) {
-        return el.isPopup;
-      }).length;
-      let numNormal = ss.getClosedWindowCount() - numPopups;
-      // #ifdef doesn't work in browser-chrome tests, so do a simple regex on platform
-      let oResults = navigator.platform.match(/Mac/) ? expectedResults.mac
-                                                     : expectedResults.other;
-      is(numPopups, oResults.popup,
-         "There were " + oResults.popup + " popup windows to repoen");
-      is(numNormal, oResults.normal,
-         "There were " + oResults.normal + " normal windows to repoen");
-
-      // cleanup & return
-      executeSoon(recCallback);
-      return;
-    }
-
-    // hack to force window to be considered a popup (toolbar=no didn't work)
-    let winData = windowsToOpen.shift();
-    let settings = "chrome,dialog=no," +
-                   (winData.isPopup ? "all=no" : "all");
-    let url = "http://example.com/?window=" + windowsToOpen.length;
-
-    provideWindow(function onTestURLLoaded(win) {
-      let tabReady = () => {
-        win.close();
-        // Give it time to close
-        executeSoon(function() {
-          openWindowRec(windowsToOpen, expectedResults, recCallback);
-        });
-      };
-
-      if (win.gMultiProcessBrowser) {
-        let tab = win.gBrowser.selectedTab;
-        tab.addEventListener("SSTabRestored", function onTabRestored() {
-          tab.removeEventListener("SSTabRestored", onTabRestored);
-          tabReady();
-        });
-      } else {
-        tabReady();
-      }
-    }, url, settings);
-  }
-
   let windowsToOpen = [{isPopup: false},
                        {isPopup: false},
                        {isPopup: true},
@@ -66,6 +67,10 @@ function test() {
                        {isPopup: true}];
   let expectedResults = {mac: {popup: 3, normal: 0},
                          other: {popup: 3, normal: 1}};
+
+  yield testWindows(windowsToOpen, expectedResults);
+
+
   let windowsToOpen2 = [{isPopup: false},
                         {isPopup: false},
                         {isPopup: false},
@@ -73,7 +78,6 @@ function test() {
                         {isPopup: false}];
   let expectedResults2 = {mac: {popup: 0, normal: 3},
                           other: {popup: 0, normal: 3}};
-  openWindowRec(windowsToOpen, expectedResults, function() {
+
-    openWindowRec(windowsToOpen2, expectedResults2, finish);
+  yield testWindows(windowsToOpen2, expectedResults2);
-  });
+});
-}

browser/components/sessionstore/test/browser_394759_perwindowpb.js

@@ -15,25 +15,27 @@ const TESTS = [
 
 function promiseTestOpenCloseWindow(aIsPrivate, aTest) {
   return Task.spawn(function*() {
-    let win = yield promiseNewWindowLoaded({ "private": aIsPrivate });
+    let win = yield BrowserTestUtils.openNewBrowserWindow({ "private": aIsPrivate });
     win.gBrowser.selectedBrowser.loadURI(aTest.url);
     yield promiseBrowserLoaded(win.gBrowser.selectedBrowser);
     yield Promise.resolve();
     // Mark the window with some unique data to be restored later on.
     ss.setWindowValue(win, aTest.key, aTest.value);
+    yield TabStateFlusher.flushWindow(win);
     // Close.
-    yield promiseWindowClosed(win);
+    yield BrowserTestUtils.closeWindow(win);
   });
 }
 
 function promiseTestOnWindow(aIsPrivate, aValue) {
   return Task.spawn(function*() {
-    let win = yield promiseNewWindowLoaded({ "private": aIsPrivate });
+    let win = yield BrowserTestUtils.openNewBrowserWindow({ "private": aIsPrivate });
+    yield TabStateFlusher.flushWindow(win);
     let data = JSON.parse(ss.getClosedWindowData())[0];
     is(ss.getClosedWindowCount(), 1, "Check that the closed window count hasn't changed");
     ok(JSON.stringify(data).indexOf(aValue) > -1,
        "Check the closed window data was stored correctly");
-    registerCleanupFunction(() => promiseWindowClosed(win));
+    registerCleanupFunction(() => BrowserTestUtils.closeWindow(win));
   });
 }
 

browser/components/sessionstore/test/browser_423132.js

@@ -66,8 +66,7 @@ function test() {
           if (gPrefService.prefHasUserValue("browser.sessionstore.interval"))
             gPrefService.clearUserPref("browser.sessionstore.interval");
           cs.removeAll();
-          newWin.close();
+          BrowserTestUtils.closeWindow(newWin).then(finish);
-          finish();
         };
 
         function flushAndReady() {

browser/components/sessionstore/test/browser_461634.js

@@ -67,8 +67,7 @@ function test() {
        "... and tabs not specifically forgetten weren't.");
 
     // clean up
-    newWin.close();
     gPrefService.clearUserPref("browser.sessionstore.max_tabs_undo");
-    finish();
+    BrowserTestUtils.closeWindow(newWin).then(finish);
   });
 }

browser/components/sessionstore/test/browser_464199.js

@@ -69,18 +69,17 @@ function test() {
 
     ForgetAboutSite.removeDataFromDomain("example.net");
     waitForClearHistory(function() {
-      closedTabs = JSON.parse(ss.getClosedTabData(newWin));
+        closedTabs = JSON.parse(ss.getClosedTabData(newWin));
-      is(closedTabs.length, remember_count,
+        is(closedTabs.length, remember_count,
-         "The correct amout of tabs was removed");
+           "The correct amout of tabs was removed");
-      is(countByTitle(closedTabs, FORGET), 0,
+        is(countByTitle(closedTabs, FORGET), 0,
-         "All tabs to be forgotten were indeed removed");
+           "All tabs to be forgotten were indeed removed");
-      is(countByTitle(closedTabs, REMEMBER), remember_count,
+        is(countByTitle(closedTabs, REMEMBER), remember_count,
-         "... and tabs to be remembered weren't.");
+           "... and tabs to be remembered weren't.");
 
-      // clean up
+        // clean up
-      newWin.close();
+        gPrefService.clearUserPref("browser.sessionstore.max_tabs_undo");
-      gPrefService.clearUserPref("browser.sessionstore.max_tabs_undo");
+        BrowserTestUtils.closeWindow(newWin).then(finish);
-     finish();
     });
   });
 }

browser/components/sessionstore/test/browser_465223.js

@@ -40,7 +40,6 @@ function test() {
        "window value was correctly overwritten");
 
     // clean up
-    newWin.close();
+    BrowserTestUtils.closeWindow(newWin).then(finish);
-    finish();
   });
 }

browser/components/sessionstore/test/browser_477657.js

@@ -52,8 +52,7 @@ function test() {
           isnot(newWin.windowState, newWin.STATE_MAXIMIZED,
                 "the window was explicitly unmaximized");
 
-          newWin.close();
+          BrowserTestUtils.closeWindow(newWin).then(finish);
-          finish();
         }, 0);
       }, 0);
     }, 0);

browser/components/sessionstore/test/browser_491577.js

@@ -114,8 +114,7 @@ function test() {
        "... and windows not specifically forgetten weren't.");
 
     // clean up
-    newWin.close();
     gPrefService.clearUserPref("browser.sessionstore.max_windows_undo");
-    finish();
+    BrowserTestUtils.closeWindow(newWin).then(finish);
   });
 }

browser/components/sessionstore/test/browser_495495.js

@@ -10,35 +10,36 @@ function test() {
   let newWin = openDialog(location, "_blank", "chrome,all,dialog=no,toolbar=yes");
   promiseWindowLoaded(newWin).then(() => {
     let state1 = ss.getWindowState(newWin);
-    newWin.close();
+    BrowserTestUtils.closeWindow(newWin).then(() => {
 
-    newWin = openDialog(location, "_blank",
+      newWin = openDialog(location, "_blank",
-      "chrome,extrachrome,menubar,resizable,scrollbars,status,toolbar=no,location,personal,directories,dialog=no");
+        "chrome,extrachrome,menubar,resizable,scrollbars,status,toolbar=no,location,personal,directories,dialog=no");
-    promiseWindowLoaded(newWin).then(() => {
+      promiseWindowLoaded(newWin).then(() => {
-      let state2 = ss.getWindowState(newWin);
+        let state2 = ss.getWindowState(newWin);
-      newWin.close();
 
-      function testState(state, expected, callback) {
+        function testState(state, expected, callback) {
-        let win = openDialog(location, "_blank", "chrome,all,dialog=no");
+          let win = openDialog(location, "_blank", "chrome,all,dialog=no");
-        promiseWindowLoaded(win).then(() => {
+          promiseWindowLoaded(win).then(() => {
 
-          is(win.gURLBar.readOnly, false,
+            is(win.gURLBar.readOnly, false,
-             "URL bar should not be read-only before setting the state");
+               "URL bar should not be read-only before setting the state");
-          is(win.gURLBar.getAttribute("enablehistory"), "true",
+            is(win.gURLBar.getAttribute("enablehistory"), "true",
-             "URL bar autocomplete should be enabled before setting the state");
+               "URL bar autocomplete should be enabled before setting the state");
-          ss.setWindowState(win, state, true);
+            ss.setWindowState(win, state, true);
-          is(win.gURLBar.readOnly, expected.readOnly,
+            is(win.gURLBar.readOnly, expected.readOnly,
-             "URL bar read-only state should be restored correctly");
+               "URL bar read-only state should be restored correctly");
-          is(win.gURLBar.getAttribute("enablehistory"), expected.enablehistory,
+            is(win.gURLBar.getAttribute("enablehistory"), expected.enablehistory,
-             "URL bar autocomplete state should be restored correctly");
+               "URL bar autocomplete state should be restored correctly");
 
-          win.close();
+            BrowserTestUtils.closeWindow(win).then(callback);
-          executeSoon(callback);
+          });
+        }
+
+        BrowserTestUtils.closeWindow(newWin).then(() => {
+          testState(state1, {readOnly: false, enablehistory: "true"}, function() {
+            testState(state2, {readOnly: true, enablehistory: "false"}, finish);
+          });
         });
-      }
-
-      testState(state1, {readOnly: false, enablehistory: "true"}, function() {
-        testState(state2, {readOnly: true, enablehistory: "false"}, finish);
       });
     });
   });

browser/components/sessionstore/test/browser_514751.js

@@ -31,8 +31,7 @@ function test() {
           gotError = true;
       }
       ok(!gotError, "Didn't get a malformed URI error.");
-      theWin.close();
+      BrowserTestUtils.closeWindow(theWin).then(finish);
-      finish();
     });
   }, false);
 }

browser/components/sessionstore/test/browser_524745.js

@@ -35,8 +35,7 @@ function test() {
 
         // Cleanup
         window.restore();
-        window_B.close();
+        BrowserTestUtils.closeWindow(window_B).then(finish);
-        finish();
       });
     }, window_B);
   });

browser/components/sessionstore/test/browser_528776.js

@@ -11,16 +11,11 @@ function browserWindowsCount(expected) {
      "number of open browser windows according to getBrowserState");
 }
 
-function test() {
+add_task(function() {
-  waitForExplicitFinish();
-
   browserWindowsCount(1);
 
-  var win = openDialog(location, "", "chrome,all,dialog=no");
+  let win = yield BrowserTestUtils.openNewBrowserWindow();
-  promiseWindowLoaded(win).then(() => {
+  browserWindowsCount(2);
-    browserWindowsCount(2);
+  yield BrowserTestUtils.closeWindow(win);
-    win.close();
+  browserWindowsCount(1);
-    browserWindowsCount(1);
+});
-    finish();
-  });
-}

browser/components/sessionstore/test/browser_580512.js

@@ -22,10 +22,17 @@ function test() {
 function closeFirstWin(win) {
   win.gBrowser.pinTab(win.gBrowser.tabs[0]);
   win.gBrowser.pinTab(win.gBrowser.tabs[1]);
+
+  let winClosed = BrowserTestUtils.windowClosed(win);
+  // We need to call BrowserTryToCloseWindow in order to trigger
+  // the machinery that chooses whether or not to save the session
+  // for the last window.
   win.BrowserTryToCloseWindow();
   ok(win.closed, "window closed");
 
-  openWinWithCb(checkSecondWin, URIS_NORMAL_B, URIS_PINNED.concat(URIS_NORMAL_B));
+  winClosed.then(() => {
+    openWinWithCb(checkSecondWin, URIS_NORMAL_B, URIS_PINNED.concat(URIS_NORMAL_B));
+  });
 }
 
 function checkSecondWin(win) {
@@ -33,11 +40,12 @@ function checkSecondWin(win) {
   is(win.gBrowser.browsers[1].currentURI.spec, URIS_PINNED[1], "second pinned tab restored");
   ok(win.gBrowser.tabs[0].pinned, "first pinned tab is still pinned");
   ok(win.gBrowser.tabs[1].pinned, "second pinned tab is still pinned");
-  win.close();
 
-  // cleanup
+  BrowserTestUtils.closeWindow(win).then(() => {
-  document.documentElement.setAttribute("windowtype", "navigator:browser");
+    // cleanup
-  finish();
+    document.documentElement.setAttribute("windowtype", "navigator:browser");
+    finish();
+  });
 }
 
 function openWinWithCb(cb, argURIs, expectedURIs) {

browser/components/sessionstore/test/browser_588426.js

@@ -10,7 +10,7 @@ function test() {
   waitForExplicitFinish();
 
   newWindowWithState(state, function (win) {
-    registerCleanupFunction(() => win.close());
+    registerCleanupFunction(() => BrowserTestUtils.closeWindow(win));
 
     is(win.gBrowser.tabs.length, 2, "two tabs were restored");
     is(win.gBrowser.visibleTabs.length, 1, "one tab is visible");

browser/components/sessionstore/test/browser_589246.js

@@ -234,8 +234,9 @@ function onWindowUnloaded() {
 
 function afterTestCleanup(aNewWin) {
   executeSoon(function() {
-    aNewWin.close();
+    BrowserTestUtils.closeWindow(aNewWin).then(() => {
-    document.documentElement.setAttribute("windowtype", originalWindowType);
+      document.documentElement.setAttribute("windowtype", originalWindowType);
-    runNextTestOrFinish();
+      runNextTestOrFinish();
+    });
   });
 }

browser/components/sessionstore/test/browser_590563.js

@@ -17,7 +17,7 @@ function test() {
   waitForExplicitFinish();
 
   newWindowWithState(state, function (win) {
-    registerCleanupFunction(() => win.close());
+    registerCleanupFunction(() => BrowserTestUtils.closeWindow(win));
 
     is(gBrowser.tabs.length, 1, "The total number of tabs should be 1");
     is(gBrowser.visibleTabs.length, 1, "The total number of visible tabs should be 1");

browser/components/sessionstore/test/browser_595601-restore_hidden.js

@@ -101,7 +101,7 @@ function newWindowWithState(state, callback) {
   let opts = "chrome,all,dialog=no,height=800,width=800";
   let win = window.openDialog(getBrowserURL(), "_blank", opts);
 
-  registerCleanupFunction(() => win.close());
+  registerCleanupFunction(() => BrowserTestUtils.closeWindow(win));
 
   whenWindowLoaded(win, function onWindowLoaded(aWin) {
     TabsProgressListener.init(aWin);

browser/components/sessionstore/test/browser_597071.js

@@ -19,16 +19,16 @@ add_task(function test_close_last_nonpopup_window() {
   ss.setWindowState(window, JSON.stringify(popupState), true);
 
   // Open a new window with a tab.
-  let win = yield promiseNewWindowLoaded({private: false});
+  let win = yield BrowserTestUtils.openNewBrowserWindow({private: false});
   let tab = win.gBrowser.addTab("http://example.com/");
-  yield promiseBrowserLoaded(tab.linkedBrowser);
+  yield BrowserTestUtils.browserLoaded(tab.linkedBrowser);
 
   // Make sure sessionstore sees this window.
   let state = JSON.parse(ss.getBrowserState());
   is(state.windows.length, 2, "sessionstore knows about this window");
 
   // Closed the window and check the closed window count.
-  yield promiseWindowClosed(win);
+  yield BrowserTestUtils.closeWindow(win);
   is(ss.getClosedWindowCount(), 1, "correct closed window count");
 
   // Cleanup.

browser/components/sessionstore/test/browser_600545.js

@@ -66,13 +66,16 @@ function done() {
   // Enumerate windows and close everything but our primary window. We can't
   // use waitForFocus() because apparently it's buggy. See bug 599253.
   let windowsEnum = Services.wm.getEnumerator("navigator:browser");
+  let closeWinPromises = [];
   while (windowsEnum.hasMoreElements()) {
     let currentWindow = windowsEnum.getNext();
     if (currentWindow != window)
-      currentWindow.close();
+      closeWinPromises.push(BrowserTestUtils.closeWindow(currentWindow));
   }
 
-  waitForBrowserState(stateBackup, finish);
+  Promise.all(closeWinPromises).then(() => {
+    waitForBrowserState(stateBackup, finish);
+  });
 }
 
 // Count up the number of tabs in the state data

browser/components/sessionstore/test/browser_615394-SSWindowState_events.js

@@ -61,15 +61,15 @@ function runNextTest() {
     // Enumerate windows and close everything but our primary window. We can't
     // use waitForFocus() because apparently it's buggy. See bug 599253.
     var windowsEnum = Services.wm.getEnumerator("navigator:browser");
+    let closeWinPromises = [];
     while (windowsEnum.hasMoreElements()) {
       var currentWindow = windowsEnum.getNext();
       if (currentWindow != window) {
-        currentWindow.close();
+        closeWinPromises.push(BrowserTestUtils.closeWindow(currentWindow));
       }
     }
 
-    // If we closed a window, give it time to close
+    Promise.all(closeWinPromises).then(() => {
-    executeSoon(function() {
       let currentTest = tests.shift();
       info("prepping for " + currentTest.name);
       waitForBrowserState(testState, currentTest);
@@ -319,9 +319,8 @@ function test_undoCloseWindow() {
 
   waitForBrowserState(lameMultiWindowState, function() {
     // Close the window which isn't window
-    newWindow.close();
+    BrowserTestUtils.closeWindow(newWindow).then(() => {
-    // Now give it time to close
+      // Now give it time to close
-    executeSoon(function() {
       reopenedWindow = ss.undoCloseWindow(0);
       reopenedWindow.addEventListener("SSWindowStateBusy", onSSWindowStateBusy, false);
       reopenedWindow.addEventListener("SSWindowStateReady", onSSWindowStateReady, false);
@@ -357,9 +356,6 @@ function test_undoCloseWindow() {
     reopenedWindow.removeEventListener("SSWindowStateReady", onSSWindowStateReady, false);
     reopenedWindow.gBrowser.tabContainer.removeEventListener("SSTabRestored", onSSTabRestored, false);
 
-    reopenedWindow.close();
+    BrowserTestUtils.closeWindow(reopenedWindow).then(runNextTest);
-
-    // Give it time to close
-    executeSoon(runNextTest);
   }
 }

browser/components/sessionstore/test/browser_618151.js

@@ -27,16 +27,19 @@ function runNextTest() {
     // Enumerate windows and close everything but our primary window. We can't
     // use waitForFocus() because apparently it's buggy. See bug 599253.
     var windowsEnum = Services.wm.getEnumerator("navigator:browser");
+    let closeWinPromises = [];
     while (windowsEnum.hasMoreElements()) {
       var currentWindow = windowsEnum.getNext();
       if (currentWindow != window) {
-        currentWindow.close();
+        closeWinPromises.push(BrowserTestUtils.closeWindow(currentWindow));
       }
     }
 
-    let currentTest = tests.shift();
+    Promise.all(closeWinPromises).then(() => {
-    info("running " + currentTest.name);
+      let currentTest = tests.shift();
-    waitForBrowserState(testState, currentTest);
+      info("running " + currentTest.name);
+      waitForBrowserState(testState, currentTest);
+    });
   }
   else {
     ss.setBrowserState(stateBackup);

browser/components/sessionstore/test/browser_625016.js

@@ -31,7 +31,7 @@ add_task(function* new_window() {
     // Double check that we have no closed windows
     is(ss.getClosedWindowCount(), 0, "no closed windows on first save");
 
-    yield promiseWindowClosed(newWin);
+    yield BrowserTestUtils.closeWindow(newWin);
     newWin = null;
 
     let state = JSON.parse((yield promiseRecoveryFileContents()));
@@ -45,7 +45,7 @@ add_task(function* new_window() {
       "observe1: 1 closed window according to API");
   } finally {
     if (newWin) {
-      yield promiseWindowClosed(newWin);
+      yield BrowserTestUtils.closeWindow(newWin);
     }
     yield forceSaveState();
   }

browser/components/sessionstore/test/browser_637020.js

@@ -55,7 +55,7 @@ add_task(function* test() {
   checkWindows();
 
   // Cleanup.
-  yield promiseWindowClosed(win);
+  yield BrowserTestUtils.closeWindow(win);
   yield promiseBrowserState(backupState);
 });
 

browser/components/sessionstore/test/browser_644409-scratchpads.js

@@ -47,8 +47,9 @@ function windowObserver(aSubject, aTopic, aData) {
             win.Scratchpad.removeObserver(this);
 
             let state = win.Scratchpad.getState();
-            win.close();
+            BrowserTestUtils.closeWindow(win).then(() => {
-            addState(state);
+              addState(state);
+            });
           },
         });
       }

browser/components/sessionstore/test/browser_659591.js

@@ -14,10 +14,11 @@ function test() {
     win.addEventListener("SSWindowClosing", function onWindowClosing() {
       win.removeEventListener("SSWindowClosing", onWindowClosing, false);
       eventReceived = true;
-      waitForFocus(finish);
     }, false);
 
-    win.close();
+    BrowserTestUtils.closeWindow(win).then(() => {
+      waitForFocus(finish);
+    });
   });
 }
 

browser/components/sessionstore/test/browser_701377.js

@@ -32,7 +32,7 @@ function newWindowWithState(aState, aCallback) {
   let opts = "chrome,all,dialog=no,height=800,width=800";
   let win = window.openDialog(getBrowserURL(), "_blank", opts);
 
-  registerCleanupFunction(() => win.close());
+  registerCleanupFunction(() => BrowserTestUtils.closeWindow(win));
 
   whenWindowLoaded(win, function onWindowLoaded(aWin) {
     ss.setWindowState(aWin, JSON.stringify(aState), true);

browser/components/sessionstore/test/browser_819510_perwindowpb.js

@@ -82,7 +82,7 @@ add_task(function* test_3() {
   is(curState.windows[2].isPrivate, true, "Window 2 is private");
   is(curState.selectedWindow, 4, "Last window opened is the one selected");
 
-  yield promiseWindowClosed(normalWindow);
+  yield BrowserTestUtils.closeWindow(normalWindow);
 
   // Pin and unpin a tab before checking the written state so that
   // the list of restoring windows gets cleared. Otherwise the

browser/components/sessionstore/test/browser_aboutSessionRestore.js

@@ -30,7 +30,7 @@ add_task(function* () {
 
   // Wait until the new window was restored.
   let win = yield waitForNewWindow();
-  yield promiseWindowClosed(win);
+  yield BrowserTestUtils.closeWindow(win);
 
   let [{tabs: [{entries: [{url}]}]}] = JSON.parse(ss.getClosedWindowData());
   is(url, "about:mozilla", "session was restored correctly");

browser/components/sessionstore/test/browser_broadcast.js

@@ -74,7 +74,7 @@ add_task(function flush_on_windowclose() {
   let browser = tab.linkedBrowser;
 
   yield modifySessionStorage(browser, {test: "on-window-close"});
-  yield closeWindow(win);
+  yield BrowserTestUtils.closeWindow(win);
 
   let [{tabs: [_, {storage}]}] = JSON.parse(ss.getClosedWindowData());
   is(storage["http://example.com"].test, "on-window-close",
@@ -138,24 +138,6 @@ function promiseNewWindow() {
   return deferred.promise;
 }
 
-function closeWindow(win) {
-  let deferred = Promise.defer();
-  let outerID = win.QueryInterface(Ci.nsIInterfaceRequestor)
-                   .getInterface(Ci.nsIDOMWindowUtils)
-                   .outerWindowID;
-
-  Services.obs.addObserver(function obs(subject, topic) {
-    let id = subject.QueryInterface(Ci.nsISupportsPRUint64).data;
-    if (id == outerID) {
-      Services.obs.removeObserver(obs, topic);
-      deferred.resolve();
-    }
-  }, "outer-window-destroyed", false);
-
-  win.close();
-  return deferred.promise;
-}
-
 function createTabWithStorageData(urls, win = window) {
   return Task.spawn(function task() {
     let tab = win.gBrowser.addTab();

browser/components/sessionstore/test/browser_cleaner.js

@@ -55,7 +55,8 @@ add_task(function* test_open_and_close() {
 
   yield promiseBrowserLoaded(tab.linkedBrowser);
 
-
+  yield TabStateFlusher.flushWindow(window);
+  yield TabStateFlusher.flushWindow(newWin);
 
   info("1. Making sure that before closing, we don't have closedAt");
   // For the moment, no "closedAt"
@@ -65,12 +66,10 @@ add_task(function* test_open_and_close() {
   is(state.windows[0].tabs[0].closedAt || false, false, "1. First tab doesn't have closedAt");
   is(state.windows[0].tabs[1].closedAt || false, false, "1. Second tab doesn't have closedAt");
 
-
-
   info("2. Making sure that after closing, we have closedAt");
 
   // Now close stuff, this should add closeAt
-  yield promiseWindowClosed(newWin);
+  yield BrowserTestUtils.closeWindow(newWin);
   yield promiseRemoveTab(newTab1);
   yield promiseRemoveTab(newTab2);
 
@@ -103,7 +102,7 @@ add_task(function* test_restore() {
   is(state.windows[0].tabs[0].closedAt || false, false, "3. First tab doesn't have closedAt");
   is(state.windows[0].tabs[1].closedAt || false, false, "3. Second tab doesn't have closedAt");
 
-  yield promiseWindowClosed(newWin);
+  yield BrowserTestUtils.closeWindow(newWin);
   gBrowser.removeTab(newTab1);
   gBrowser.removeTab(newTab2);
 });

browser/components/sessionstore/test/browser_dying_cache.js

@@ -31,7 +31,7 @@ add_task(function* test() {
   let closedTabData = ss.getClosedTabData(win);
 
   // Close our window.
-  yield promiseWindowClosed(win);
+  yield BrowserTestUtils.closeWindow(win);
 
   // SessionStore should no longer track our window
   // but it should still report the same state.

browser/components/sessionstore/test/browser_merge_closed_tabs.js

@@ -65,7 +65,7 @@ add_task(function () {
 
   // Clean up.
   gPrefService.clearUserPref("browser.sessionstore.max_tabs_undo");
-  win.close();
+  yield BrowserTestUtils.closeWindow(win);
 });
 
 

browser/components/sessionstore/test/browser_privatetabs.js

@@ -97,7 +97,7 @@ add_task(function () {
 
   // Check that all private tabs are removed when the non-private
   // window is closed and we don't save windows without any tabs.
-  yield promiseWindowClosed(win);
+  yield BrowserTestUtils.closeWindow(win);
   is(ss.getClosedWindowCount(), 0, "no windows to restore");
 });
 
@@ -123,7 +123,7 @@ add_task(function () {
   is(ss.getClosedTabCount(win), 1, "there is a single tab to restore");
 
   // Ensure that closed private windows can never be restored.
-  yield promiseWindowClosed(win);
+  yield BrowserTestUtils.closeWindow(win);
   is(ss.getClosedWindowCount(), 0, "no windows to restore");
 });
 

browser/components/sessionstore/test/browser_swapDocShells.js

@@ -14,7 +14,7 @@ add_task(function* () {
   let browser = win.gBrowser.browsers[1];
   is(browser.currentURI.spec, "about:mozilla", "tab was duplicated");
 
-  yield promiseWindowClosed(win);
+  yield BrowserTestUtils.closeWindow(win);
 });
 
 function promiseDelayedStartupFinished(win) {

browser/components/sessionstore/test/browser_windowRestore_perwindowpb.js

@@ -17,11 +17,10 @@ function test() {
     win.addEventListener("SSWindowClosing", function onclosing() {
       win.removeEventListener("SSWindowClosing", onclosing, false);
       executeSoon(function () {
-        is (ss.getClosedWindowCount(), 0,
+        is(ss.getClosedWindowCount(), 0,
             "The private window should not have been stored");
-        finish();
       });
     }, false);
-    win.close();
+    BrowserTestUtils.closeWindow(win).then(finish);
   });
 }

browser/components/sessionstore/test/head.js

@@ -400,7 +400,7 @@ function promiseAllButPrimaryWindowClosed() {
     }
   }
 
-  return Promise.all(windows.map(promiseWindowClosed));
+  return Promise.all(windows.map(BrowserTestUtils.closeWindow));
 }
 
 // Forget all closed windows.
@@ -438,7 +438,8 @@ function whenNewWindowLoaded(aOptions, aCallback) {
 
     win.addEventListener("load", function onLoad() {
       win.removeEventListener("load", onLoad);
-      resolve(promiseBrowserLoaded(win.gBrowser.selectedBrowser));
+      let browser = win.gBrowser.selectedBrowser;
+      promiseBrowserLoaded(browser).then(resolve);
     });
   });
 
@@ -448,24 +449,6 @@ function promiseNewWindowLoaded(aOptions) {
   return new Promise(resolve => whenNewWindowLoaded(aOptions, resolve));
 }
 
-/**
- * Chrome windows aren't closed synchronously. Provide a helper method to close
- * a window and wait until we received the "domwindowclosed" notification for it.
- */
-function promiseWindowClosed(win) {
-  let promise = new Promise(resolve => {
-    Services.obs.addObserver(function obs(subject, topic) {
-      if (subject == win) {
-        Services.obs.removeObserver(obs, topic);
-        resolve();
-      }
-    }, "domwindowclosed", false);
-  });
-
-  win.close();
-  return promise;
-}
-
 function runInContent(browser, func, arg, callback = null) {
   let deferred = Promise.defer();
 

testing/mochitest/BrowserTestUtils/BrowserTestUtils.jsm

@@ -318,10 +318,12 @@ this.BrowserTestUtils = {
    *        A window to close.
    *
    * @return {Promise}
-   *         Resolves when the provided window has been closed.
+   *         Resolves when the provided window has been closed. For browser
+   *         windows, the Promise will also wait until all final SessionStore
+   *         messages have been sent up from all browser tabs.
    */
   closeWindow(win) {
-    return new Promise(resolve => {
+    let domWinClosedPromise = new Promise((resolve) => {
       function observer(subject, topic, data) {
         if (topic == "domwindowclosed" && subject === win) {
           Services.ww.unregisterNotification(observer);
@@ -329,8 +331,36 @@ this.BrowserTestUtils = {
         }
       }
       Services.ww.registerNotification(observer);
-      win.close();
     });
+
+    let promises = [domWinClosedPromise];
+    let winType = win.document.documentElement.getAttribute("windowtype");
+
+    if (winType == "navigator:browser") {
+      let finalMsgsPromise = new Promise((resolve) => {
+        let browserSet = new Set(win.gBrowser.browsers);
+        let mm = win.getGroupMessageManager("browsers");
+
+        mm.addMessageListener("SessionStore:update", function onMessage(msg) {
+          if (browserSet.has(msg.target) && msg.data.isFinal) {
+            browserSet.delete(msg.target);
+            if (!browserSet.size) {
+              mm.removeMessageListener("SessionStore:update", onMessage);
+              // Give the TabStateFlusher a chance to react to this final
+              // update and for the TabStateFlusher.flushWindow promise
+              // to resolve before we resolve.
+              resolve();
+            }
+          }
+        }, true);
+      });
+
+      promises.push(finalMsgsPromise);
+    }
+
+    win.close();
+
+    return Promise.all(promises);
   },
 
   /**