Bug 1171708 - Fix SessionStore tests to account for async window flushing. r?billm draft
authorMike Conley <mconley@mozilla.com>
Fri, 06 Nov 2015 17:01:15 -0500
changeset 308846 e3cd17601fc8f58d8161e61459012235cf520222
parent 308845 d674ea45bce3b413dabfabfc9cf62e92f29ee0c1
child 308847 52d9d917e4a56d50ee5c60320f9f60bd0c3ae93a
push id7541
push usermconley@mozilla.com
push dateMon, 16 Nov 2015 00:50:56 +0000
reviewersbillm
bugs1171708
milestone45.0a1
Bug 1171708 - Fix SessionStore tests to account for async window flushing. r?billm
browser/components/sessionstore/test/browser_394759_behavior.js
browser/components/sessionstore/test/browser_394759_perwindowpb.js
browser/components/sessionstore/test/browser_464199.js
browser/components/sessionstore/test/browser_528776.js
browser/components/sessionstore/test/browser_595601-restore_hidden.js
browser/components/sessionstore/test/browser_597071.js
browser/components/sessionstore/test/browser_615394-SSWindowState_events.js
browser/components/sessionstore/test/browser_625016.js
browser/components/sessionstore/test/browser_637020.js
browser/components/sessionstore/test/browser_819510_perwindowpb.js
browser/components/sessionstore/test/browser_aboutSessionRestore.js
browser/components/sessionstore/test/browser_broadcast.js
browser/components/sessionstore/test/browser_cleaner.js
browser/components/sessionstore/test/browser_dying_cache.js
browser/components/sessionstore/test/browser_privatetabs.js
browser/components/sessionstore/test/browser_swapDocShells.js
browser/components/sessionstore/test/head.js
testing/mochitest/BrowserTestUtils/BrowserTestUtils.jsm
--- a/browser/components/sessionstore/test/browser_394759_behavior.js
+++ b/browser/components/sessionstore/test/browser_394759_behavior.js
@@ -1,79 +1,83 @@
-/* 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
- * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+/**
+ * Test helper function that opens a series of windows, closes them
+ * 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},
                        {isPopup: true},
                        {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},
                         {isPopup: false},
                         {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);
+});
\ No newline at end of file
--- a/browser/components/sessionstore/test/browser_394759_perwindowpb.js
+++ b/browser/components/sessionstore/test/browser_394759_perwindowpb.js
@@ -10,35 +10,37 @@ const TESTS = [
     value: "uniq" + r() },
   { url: "about:mozilla",
     key: "bug 394759 PB",
     value: "uniq" + r() },
 ];
 
 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));
   });
 }
 
 add_task(function* init() {
   forgetClosedWindows();
   while (ss.getClosedTabCount(window) > 0) {
     ss.forgetClosedTab(window, 0);
   }
--- a/browser/components/sessionstore/test/browser_464199.js
+++ b/browser/components/sessionstore/test/browser_464199.js
@@ -64,23 +64,23 @@ function test() {
     is(countByTitle(closedTabs, FORGET),
        test_state.windows[0]._closedTabs.length - remember_count,
        "The correct amout of tabs are to be forgotten");
     is(countByTitle(closedTabs, REMEMBER), remember_count,
        "Everything is set up.");
 
     ForgetAboutSite.removeDataFromDomain("example.net");
     waitForClearHistory(function() {
-      closedTabs = JSON.parse(ss.getClosedTabData(newWin));
-      is(closedTabs.length, remember_count,
-         "The correct amout of tabs was removed");
-      is(countByTitle(closedTabs, FORGET), 0,
-         "All tabs to be forgotten were indeed removed");
-      is(countByTitle(closedTabs, REMEMBER), remember_count,
-         "... and tabs to be remembered weren't.");
+        closedTabs = JSON.parse(ss.getClosedTabData(newWin));
+        is(closedTabs.length, remember_count,
+           "The correct amout of tabs was removed");
+        is(countByTitle(closedTabs, FORGET), 0,
+           "All tabs to be forgotten were indeed removed");
+        is(countByTitle(closedTabs, REMEMBER), remember_count,
+           "... and tabs to be remembered weren't.");
 
-      // clean up
-      newWin.close();
-      gPrefService.clearUserPref("browser.sessionstore.max_tabs_undo");
-     finish();
+        // clean up
+        newWin.close();
+        gPrefService.clearUserPref("browser.sessionstore.max_tabs_undo");
+        finish();
     });
   });
 }
--- a/browser/components/sessionstore/test/browser_528776.js
+++ b/browser/components/sessionstore/test/browser_528776.js
@@ -6,21 +6,16 @@ function browserWindowsCount(expected) {
       ++count;
   }
   is(count, expected,
      "number of open browser windows according to nsIWindowMediator");
   is(JSON.parse(ss.getBrowserState()).windows.length, expected,
      "number of open browser windows according to getBrowserState");
 }
 
-function test() {
-  waitForExplicitFinish();
-
+add_task(function() {
   browserWindowsCount(1);
 
-  var win = openDialog(location, "", "chrome,all,dialog=no");
-  promiseWindowLoaded(win).then(() => {
-    browserWindowsCount(2);
-    win.close();
-    browserWindowsCount(1);
-    finish();
-  });
-}
+  let win = yield BrowserTestUtils.openNewBrowserWindow();
+  browserWindowsCount(2);
+  yield BrowserTestUtils.closeWindow(win);
+  browserWindowsCount(1);
+});
--- a/browser/components/sessionstore/test/browser_595601-restore_hidden.js
+++ b/browser/components/sessionstore/test/browser_595601-restore_hidden.js
@@ -96,17 +96,17 @@ var TabsProgressListener = {
   }
 }
 
 // ----------
 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);
     TabsProgressListener.setCallback(callback);
 
     ss.setWindowState(aWin, JSON.stringify(state), true);
   });
 }
--- a/browser/components/sessionstore/test/browser_597071.js
+++ b/browser/components/sessionstore/test/browser_597071.js
@@ -14,23 +14,23 @@ add_task(function test_close_last_nonpop
   let popupState = {windows: [
     {tabs: [{entries: []}], isPopup: true, hidden: "toolbar"}
   ]};
 
   // Set this window to be a popup.
   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.
   ss.setWindowState(window, oldState, true);
 });
--- a/browser/components/sessionstore/test/browser_615394-SSWindowState_events.js
+++ b/browser/components/sessionstore/test/browser_615394-SSWindowState_events.js
@@ -314,19 +314,18 @@ function test_undoCloseWindow() {
       newWindow = aSubject.QueryInterface(Ci.nsIDOMWindow);
       Services.ww.unregisterNotification(firstWindowObserver);
     }
   }
   Services.ww.registerNotification(firstWindowObserver);
 
   waitForBrowserState(lameMultiWindowState, function() {
     // Close the window which isn't window
-    newWindow.close();
-    // Now give it time to close
-    executeSoon(function() {
+    BrowserTestUtils.closeWindow(newWindow).then(() => {
+      // Now give it time to close
       reopenedWindow = ss.undoCloseWindow(0);
       reopenedWindow.addEventListener("SSWindowStateBusy", onSSWindowStateBusy, false);
       reopenedWindow.addEventListener("SSWindowStateReady", onSSWindowStateReady, false);
 
       reopenedWindow.addEventListener("load", function() {
         reopenedWindow.removeEventListener("load", arguments.callee, false);
 
         reopenedWindow.gBrowser.tabContainer.addEventListener("SSTabRestored", onSSTabRestored, false);
--- a/browser/components/sessionstore/test/browser_625016.js
+++ b/browser/components/sessionstore/test/browser_625016.js
@@ -26,31 +26,31 @@ add_task(function* new_window() {
   try {
     newWin = yield promiseNewWindowLoaded();
     let tab = newWin.gBrowser.addTab("http://example.com/browser_625016.js?" + Math.random());
     yield promiseBrowserLoaded(tab.linkedBrowser);
 
     // 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()));
     is(state.windows.length, 2,
       "observe1: 2 windows in data written to disk");
     is(state._closedWindows.length, 0,
       "observe1: no closed windows in data written to disk");
 
     // The API still treats the closed window as closed, so ensure that window is there
     is(ss.getClosedWindowCount(), 1,
       "observe1: 1 closed window according to API");
   } finally {
     if (newWin) {
-      yield promiseWindowClosed(newWin);
+      yield BrowserTestUtils.closeWindow(newWin);
     }
     yield forceSaveState();
   }
 });
 
 // We'll open a tab, which should trigger another state save which would wipe
 // the _shouldRestore attribute from the closed window
 add_task(function* new_tab() {
--- a/browser/components/sessionstore/test/browser_637020.js
+++ b/browser/components/sessionstore/test/browser_637020.js
@@ -50,17 +50,17 @@ add_task(function* test() {
 
   // The history has now been restored and the tabs are loading. The data must
   // now come from the window, if it's correctly been marked as dirty before.
   yield new Promise(resolve => whenDelayedStartupFinished(win, resolve));
   info("the delayed startup has finished");
   checkWindows();
 
   // Cleanup.
-  yield promiseWindowClosed(win);
+  yield BrowserTestUtils.closeWindow(win);
   yield promiseBrowserState(backupState);
 });
 
 function checkWindows() {
   let state = JSON.parse(SessionStore.getBrowserState());
   is(state.windows[0].tabs.length, 2, "first window has two tabs");
   is(state.windows[1].tabs.length, 2, "second window has two tabs");
 }
--- a/browser/components/sessionstore/test/browser_819510_perwindowpb.js
+++ b/browser/components/sessionstore/test/browser_819510_perwindowpb.js
@@ -77,17 +77,17 @@ add_task(function* test_3() {
   win = yield promiseNewWindowLoaded();
   yield promiseTabLoad(win, "http://www.example.com/");
 
   let curState = JSON.parse(ss.getBrowserState());
   is(curState.windows.length, 4, "Browser has opened 4 windows");
   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
   // window we just closed would be marked as not closed.
   let tab = win.gBrowser.tabs[0];
   win.gBrowser.pinTab(tab);
   win.gBrowser.unpinTab(tab);
 
--- a/browser/components/sessionstore/test/browser_aboutSessionRestore.js
+++ b/browser/components/sessionstore/test/browser_aboutSessionRestore.js
@@ -25,17 +25,17 @@ add_task(function* () {
   ss.setTabState(tab, JSON.stringify(TAB_STATE));
   yield promiseTabRestored(tab);
 
   ok(gBrowser.tabs.length > 1, "we have more than one tab");
   browser.messageManager.loadFrameScript(FRAME_SCRIPT, true);
 
   // 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");
   ss.forgetClosedWindow(0);
 });
 
 function waitForNewWindow() {
   return new Promise(resolve => {
--- a/browser/components/sessionstore/test/browser_broadcast.js
+++ b/browser/components/sessionstore/test/browser_broadcast.js
@@ -69,17 +69,17 @@ add_task(function flush_on_duplicate() {
  * a window is closed.
  */
 add_task(function flush_on_windowclose() {
   let win = yield promiseNewWindow();
   let tab = yield createTabWithStorageData(["http://example.com"], win);
   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",
     "sessionStorage data has been flushed when closing a window");
 });
 
 /**
  * This test ensures that stale tab data is ignored when reusing a tab
@@ -133,34 +133,16 @@ add_task(function flush_on_tabclose_racy
 });
 
 function promiseNewWindow() {
   let deferred = Promise.defer();
   whenNewWindowLoaded({private: false}, deferred.resolve);
   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();
     let browser = tab.linkedBrowser;
 
     for (let url of urls) {
       browser.loadURI(url);
       yield promiseBrowserLoaded(browser);
--- a/browser/components/sessionstore/test/browser_cleaner.js
+++ b/browser/components/sessionstore/test/browser_cleaner.js
@@ -50,32 +50,31 @@ add_task(function* test_open_and_close()
   let newTab2 = gBrowser.addTab(URL_TAB2);
   yield promiseBrowserLoaded(newTab2.linkedBrowser);
 
   let newWin = yield promiseNewWindowLoaded();
   let tab = newWin.gBrowser.addTab(URL_NEWWIN);
 
   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"
   let state = JSON.parse(ss.getBrowserState());
   is(state.windows[0].closedAt || false, false, "1. Main window doesn't have closedAt");
   is(state.windows[1].closedAt || false, false, "1. Second window doesn't have closedAt");
   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);
 
   state = CLOSED_STATE = JSON.parse(ss.getBrowserState());
 
   is(state.windows[0].closedAt || false, false, "2. Main window doesn't have closedAt");
   ok(isRecent(state._closedWindows[0].closedAt), "2. Second window was closed recently");
   ok(isRecent(state.windows[0]._closedTabs[0].closedAt), "2. First tab was closed recently");
@@ -98,17 +97,17 @@ add_task(function* test_restore() {
 
   let state = JSON.parse(ss.getBrowserState());
 
   is(state.windows[0].closedAt || false, false, "3. Main window doesn't have closedAt");
   is(state.windows[1].closedAt || false, false, "3. Second window doesn't have closedAt");
   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);
 });
 
 
 add_task(function* test_old_data() {
   info("4. Removing closedAt from the sessionstore, making sure that it is added upon idle-daily");
 
--- a/browser/components/sessionstore/test/browser_dying_cache.js
+++ b/browser/components/sessionstore/test/browser_dying_cache.js
@@ -26,17 +26,17 @@ add_task(function* test() {
   ok("__SSi" in win, "window is being tracked by sessionstore");
   ss.setWindowValue(win, "foo", "bar");
   checkWindowState(win);
 
   let state = ss.getWindowState(win);
   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.
   ok(!("__SSi" in win), "sessionstore does no longer track our window");
   checkWindowState(win);
 
   // Make sure we're not allowed to modify state data.
   Assert.throws(() => ss.setWindowState(win, {}),
--- a/browser/components/sessionstore/test/browser_privatetabs.js
+++ b/browser/components/sessionstore/test/browser_privatetabs.js
@@ -92,17 +92,17 @@ add_task(function () {
   yield TabStateFlusher.flush(browser);
 
   // Check that we consider the tab as private.
   state = JSON.parse(ss.getTabState(tab));
   ok(state.isPrivate, "tab considered private");
 
   // 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");
 });
 
 add_task(function () {
   // Clear the list of closed windows.
   forgetClosedWindows();
 
   // Create a new window to attach our frame script to.
@@ -118,16 +118,16 @@ add_task(function () {
   let state = JSON.parse(ss.getTabState(tab));
   ok(state.isPrivate, "tab considered private");
 
   // Ensure that closed tabs in a private windows can be restored.
   win.gBrowser.removeTab(tab);
   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");
 });
 
 function setUsePrivateBrowsing(browser, val) {
   return sendMessage(browser, "ss-test:setUsePrivateBrowsing", val);
 }
 
--- a/browser/components/sessionstore/test/browser_swapDocShells.js
+++ b/browser/components/sessionstore/test/browser_swapDocShells.js
@@ -9,17 +9,17 @@ add_task(function* () {
   yield promiseBrowserHasURL(win.gBrowser.browsers[0], "about:mozilla");
 
   win.duplicateTabIn(win.gBrowser.selectedTab, "tab");
   yield promiseTabRestored(win.gBrowser.tabs[1]);
 
   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) {
   let deferred = Promise.defer();
   whenDelayedStartupFinished(win, deferred.resolve);
   return deferred.promise;
 }
 
--- a/browser/components/sessionstore/test/head.js
+++ b/browser/components/sessionstore/test/head.js
@@ -395,17 +395,17 @@ registerCleanupFunction(function () {
 function promiseAllButPrimaryWindowClosed() {
   let windows = [];
   for (let win in BrowserWindowIterator()) {
     if (win != window) {
       windows.push(win);
     }
   }
 
-  return Promise.all(windows.map(promiseWindowClosed));
+  return Promise.all(windows.map(BrowserTestUtils.closeWindow));
 }
 
 // Forget all closed windows.
 function forgetClosedWindows() {
   while (ss.getClosedWindowCount() > 0) {
     ss.forgetClosedWindow(0);
   }
 }
@@ -443,34 +443,16 @@ function whenNewWindowLoaded(aOptions, a
   });
 
   Promise.all([delayedStartup, browserLoaded]).then(() => aCallback(win));
 }
 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();
 
   let mm = browser.messageManager;
   mm.sendAsyncMessage("ss-test:run", {code: func.toSource()}, {arg: arg});
   mm.addMessageListener("ss-test:runFinished", ({data}) => deferred.resolve(data));
 
   return deferred.promise;
--- a/testing/mochitest/BrowserTestUtils/BrowserTestUtils.jsm
+++ b/testing/mochitest/BrowserTestUtils/BrowserTestUtils.jsm
@@ -313,29 +313,56 @@ this.BrowserTestUtils = {
 
   /**
    * Closes a window.
    *
    * @param {Window}
    *        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);
           resolve();
         }
       }
       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);
+              resolve();
+            }
+          }
+        }, true);
+      });
+
+      promises.push(finalMsgsPromise);
+    }
+
+    win.close();
+
+    return Promise.all(promises);
   },
 
   /**
    * Waits for an event to be fired on a specified element.
    *
    * Usage:
    *    let promiseEvent = BrowserTestUtil.waitForEvent(element, "eventName");
    *    // Do some processing here that will cause the event to be fired