Bug 1034036 - Part 6: Tests that use ss.setBrowserState() or ss.setWindowState() should wait until the window is restored to continue. r=dao

MozReview-Commit-ID: 5d1E5TjKnIR

--HG--
extra : rebase_source : be521af44edeae3f218602b4e279cfeb17ae866f

browser/components/extensions/test/browser/browser_ext_tabs_lazy.js

@@ -14,6 +14,7 @@ const SESSION = {
 
 add_task(async function() {
   SessionStore.setBrowserState(JSON.stringify(SESSION));
+  await promiseWindowRestored(window);
   const tab = gBrowser.tabs[1];
 
   is(tab.getAttribute("pending"), "true", "The tab is pending restore");

browser/components/extensions/test/browser/head.js

@@ -21,7 +21,7 @@
  *          promisePrefChangeObserved openContextMenuInFrame
  *          promiseAnimationFrame getCustomizableUIPanelID
  *          awaitEvent BrowserWindowIterator
- *          navigateTab historyPushState
+ *          navigateTab historyPushState promiseWindowRestored
  */
 
 // There are shutdown issues for which multiple rejections are left uncaught.
@@ -488,6 +488,10 @@ function promisePrefChangeObserved(pref) {
     }));
 }
 
+function promiseWindowRestored(window) {
+  return new Promise(resolve => window.addEventListener("SSWindowRestored", resolve, {once: true}));
+}
+
 function awaitEvent(eventName, id) {
   return new Promise(resolve => {
     let listener = (_eventName, ...args) => {

browser/components/sessionstore/test/browser_354894_perwindowpb.js

@@ -439,8 +439,6 @@ add_task(async function test_open_close_restore_from_popup() {
 
     newWin2 = await promiseNewWindowLoaded();
 
-    is(newWin2.gBrowser.browsers.length, 1,
-       "Did not restore, as undoCloseWindow() was last called");
     is(TEST_URLS.indexOf(newWin2.gBrowser.browsers[0].currentURI.spec), -1,
        "Did not restore, as undoCloseWindow() was last called (2)");
 

browser/components/sessionstore/test/browser_423132.js

@@ -36,7 +36,7 @@ add_task(async function() {
   Services.cookies.removeAll();
 
   // restore the window state
-  ss.setBrowserState(state);
+  await setBrowserState(state);
 
   // at this point, the cookie should be restored...
   enumerator = Services.cookies.enumerator;

browser/components/sessionstore/test/browser_461634.js

@@ -4,11 +4,9 @@
 
 ChromeUtils.import("resource:///modules/sessionstore/SessionStore.jsm");
 
-function test() {
+add_task(async function testClosedTabData() {
   /** Test for Bug 461634 **/
 
-  waitForExplicitFinish();
-
   const REMEMBER = Date.now(), FORGET = Math.random();
   let test_state = { windows: [{ "tabs": [{ "entries": [] }], _closedTabs: [
     { state: { entries: [{ url: "http://www.example.net/" }] }, title: FORGET },
@@ -33,52 +31,52 @@ function test() {
 
   // Open a window and add the above closed tab list.
   let newWin = openDialog(location, "", "chrome,all,dialog=no");
-  promiseWindowLoaded(newWin).then(() => {
-    Services.prefs.setIntPref("browser.sessionstore.max_tabs_undo",
-                              test_state.windows[0]._closedTabs.length);
-    ss.setWindowState(newWin, JSON.stringify(test_state), true);
+  await promiseWindowLoaded(newWin);
 
-    let closedTabs = SessionStore.getClosedTabData(newWin, false);
+  Services.prefs.setIntPref("browser.sessionstore.max_tabs_undo",
+                            test_state.windows[0]._closedTabs.length);
+  await setWindowState(newWin, test_state);
 
-    // Verify that non JSON serialized data is the same as JSON serialized data.
-    is(JSON.stringify(closedTabs), SessionStore.getClosedTabData(newWin),
-       "Non-serialized data is the same as serialized data");
+  let closedTabs = SessionStore.getClosedTabData(newWin, false);
 
-    is(closedTabs.length, test_state.windows[0]._closedTabs.length,
-       "Closed tab list has the expected length");
-    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");
+  // Verify that non JSON serialized data is the same as JSON serialized data.
+  is(JSON.stringify(closedTabs), SessionStore.getClosedTabData(newWin),
+     "Non-serialized data is the same as serialized data");
 
-    // All of the following calls with illegal arguments should throw NS_ERROR_ILLEGAL_VALUE.
-    ok(testForError(() => ss.forgetClosedTab({}, 0)),
-       "Invalid window for forgetClosedTab throws");
-    ok(testForError(() => ss.forgetClosedTab(newWin, -1)),
-       "Invalid tab for forgetClosedTab throws");
-    ok(testForError(() => ss.forgetClosedTab(newWin, test_state.windows[0]._closedTabs.length + 1)),
-       "Invalid tab for forgetClosedTab throws");
+  is(closedTabs.length, test_state.windows[0]._closedTabs.length,
+     "Closed tab list has the expected length");
+  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");
 
-    // Remove third tab, then first tab.
-    ss.forgetClosedTab(newWin, 2);
-    ss.forgetClosedTab(newWin, null);
+  // All of the following calls with illegal arguments should throw NS_ERROR_ILLEGAL_VALUE.
+  ok(testForError(() => ss.forgetClosedTab({}, 0)),
+     "Invalid window for forgetClosedTab throws");
+  ok(testForError(() => ss.forgetClosedTab(newWin, -1)),
+     "Invalid tab for forgetClosedTab throws");
+  ok(testForError(() => ss.forgetClosedTab(newWin, test_state.windows[0]._closedTabs.length + 1)),
+     "Invalid tab for forgetClosedTab throws");
 
-    closedTabs = SessionStore.getClosedTabData(newWin, false);
+  // Remove third tab, then first tab.
+  ss.forgetClosedTab(newWin, 2);
+  ss.forgetClosedTab(newWin, null);
 
-    // Verify that non JSON serialized data is the same as JSON serialized data.
-    is(JSON.stringify(closedTabs), SessionStore.getClosedTabData(newWin),
-       "Non-serialized data is the same as serialized data");
+  closedTabs = SessionStore.getClosedTabData(newWin, false);
 
-    is(closedTabs.length, remember_count,
-       "The correct amout of tabs was removed");
-    is(countByTitle(closedTabs, FORGET), 0,
-       "All tabs specifically forgotten were indeed removed");
-    is(countByTitle(closedTabs, REMEMBER), remember_count,
-       "... and tabs not specifically forgetten weren't");
+  // Verify that non JSON serialized data is the same as JSON serialized data.
+  is(JSON.stringify(closedTabs), SessionStore.getClosedTabData(newWin),
+     "Non-serialized data is the same as serialized data");
 
-    // Clean up.
-    Services.prefs.clearUserPref("browser.sessionstore.max_tabs_undo");
-    BrowserTestUtils.closeWindow(newWin).then(finish);
-  });
-}
+  is(closedTabs.length, remember_count,
+     "The correct amout of tabs was removed");
+  is(countByTitle(closedTabs, FORGET), 0,
+     "All tabs specifically forgotten were indeed removed");
+  is(countByTitle(closedTabs, REMEMBER), remember_count,
+     "... and tabs not specifically forgetten weren't");
+
+  // Clean up.
+  Services.prefs.clearUserPref("browser.sessionstore.max_tabs_undo");
+  await BrowserTestUtils.closeWindow(newWin);
+});

browser/components/sessionstore/test/browser_464199.js

@@ -57,6 +57,7 @@ add_task(async function() {
   Services.prefs.setIntPref("browser.sessionstore.max_tabs_undo",
                             test_state.windows[0]._closedTabs.length);
   ss.setWindowState(newWin, JSON.stringify(test_state), true);
+  await promiseWindowRestored(newWin);
 
   let closedTabs = JSON.parse(ss.getClosedTabData(newWin));
   is(closedTabs.length, test_state.windows[0]._closedTabs.length,

browser/components/sessionstore/test/browser_465223.js

@@ -2,12 +2,9 @@
  * 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/. */
 
-function test() {
+add_task(async function test_clearWindowValues() {
   /** Test for Bug 465223 **/
 
-  // test setup
-  waitForExplicitFinish();
-
   let uniqueKey1 = "bug 465223.1";
   let uniqueKey2 = "bug 465223.2";
   let uniqueValue1 = "unik" + Date.now();
@@ -15,31 +12,30 @@ function test() {
 
   // open a window and set a value on it
   let newWin = openDialog(location, "_blank", "chrome,all,dialog=no");
-  promiseWindowLoaded(newWin).then(() => {
-    ss.setWindowValue(newWin, uniqueKey1, uniqueValue1);
+  await promiseWindowLoaded(newWin);
+  ss.setWindowValue(newWin, uniqueKey1, uniqueValue1);
 
-    let newState = { windows: [{ tabs: [{ entries: [] }], extData: {} }] };
-    newState.windows[0].extData[uniqueKey2] = uniqueValue2;
-    ss.setWindowState(newWin, JSON.stringify(newState), false);
+  let newState = { windows: [{ tabs: [{ entries: [] }], extData: {} }] };
+  newState.windows[0].extData[uniqueKey2] = uniqueValue2;
+  await setWindowState(newWin, newState);
 
-    is(newWin.gBrowser.tabs.length, 2,
-       "original tab wasn't overwritten");
-    is(ss.getWindowValue(newWin, uniqueKey1), uniqueValue1,
-       "window value wasn't overwritten when the tabs weren't");
-    is(ss.getWindowValue(newWin, uniqueKey2), uniqueValue2,
-       "new window value was correctly added");
+  is(newWin.gBrowser.tabs.length, 2,
+    "original tab wasn't overwritten");
+  is(ss.getWindowValue(newWin, uniqueKey1), uniqueValue1,
+    "window value wasn't overwritten when the tabs weren't");
+  is(ss.getWindowValue(newWin, uniqueKey2), uniqueValue2,
+    "new window value was correctly added");
 
-    newState.windows[0].extData[uniqueKey2] = uniqueValue1;
-    ss.setWindowState(newWin, JSON.stringify(newState), true);
+  newState.windows[0].extData[uniqueKey2] = uniqueValue1;
+  await setWindowState(newWin, newState, true);
 
-    is(newWin.gBrowser.tabs.length, 1,
-       "original tabs were overwritten");
-    is(ss.getWindowValue(newWin, uniqueKey1), "",
-       "window value was cleared");
-    is(ss.getWindowValue(newWin, uniqueKey2), uniqueValue1,
-       "window value was correctly overwritten");
+  is(newWin.gBrowser.tabs.length, 1,
+    "original tabs were overwritten");
+  is(ss.getWindowValue(newWin, uniqueKey1), "",
+    "window value was cleared");
+  is(ss.getWindowValue(newWin, uniqueKey2), uniqueValue1,
+    "window value was correctly overwritten");
 
-    // clean up
-    BrowserTestUtils.closeWindow(newWin).then(finish);
-  });
-}
+  // clean up
+  await BrowserTestUtils.closeWindow(newWin);
+});

browser/components/sessionstore/test/browser_477657.js

@@ -2,59 +2,56 @@
  * 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/. */
 
-function test() {
+add_task(async function test_sizemodeDefaults() {
   /** Test for Bug 477657 **/
-  waitForExplicitFinish();
-
   let newWin = openDialog(location, "_blank", "chrome,all,dialog=no");
-  promiseWindowLoaded(newWin).then(() => {
-    let newState = { windows: [{
-      tabs: [{ entries: [] }],
-      _closedTabs: [{
-        state: { entries: [{ url: "about:" }]},
-        title: "About:"
-      }],
-      sizemode: "maximized"
-    }] };
+  await promiseWindowLoaded(newWin);
+  let newState = { windows: [{
+    tabs: [{ entries: [] }],
+    _closedTabs: [{
+      state: { entries: [{ url: "about:" }]},
+      title: "About:"
+    }],
+    sizemode: "maximized"
+  }] };
 
-    let uniqueKey = "bug 477657";
-    let uniqueValue = "unik" + Date.now();
+  let uniqueKey = "bug 477657";
+  let uniqueValue = "unik" + Date.now();
 
-    ss.setWindowValue(newWin, uniqueKey, uniqueValue);
-    is(ss.getWindowValue(newWin, uniqueKey), uniqueValue,
-       "window value was set before the window was overwritten");
-    ss.setWindowState(newWin, JSON.stringify(newState), true);
+  ss.setWindowValue(newWin, uniqueKey, uniqueValue);
+  is(ss.getWindowValue(newWin, uniqueKey), uniqueValue,
+     "window value was set before the window was overwritten");
 
-    // use newWin.setTimeout(..., 0) to mirror sss_restoreWindowFeatures
-    newWin.setTimeout(function() {
-      is(ss.getWindowValue(newWin, uniqueKey), "",
-         "window value was implicitly cleared");
+  await setWindowState(newWin, newState, true);
+  // use newWin.setTimeout(..., 0) to mirror sss_restoreWindowFeatures
+  await new Promise(resolve => newWin.setTimeout(resolve, 0));
 
-      is(newWin.windowState, newWin.STATE_MAXIMIZED,
-         "the window was maximized");
+  is(ss.getWindowValue(newWin, uniqueKey), "",
+    "window value was implicitly cleared");
 
-      is(JSON.parse(ss.getClosedTabData(newWin)).length, 1,
-         "the closed tab was added before the window was overwritten");
-      delete newState.windows[0]._closedTabs;
-      delete newState.windows[0].sizemode;
-      ss.setWindowState(newWin, JSON.stringify(newState), true);
+  is(newWin.windowState, newWin.STATE_MAXIMIZED,
+    "the window was maximized");
 
-      newWin.setTimeout(function() {
-        is(JSON.parse(ss.getClosedTabData(newWin)).length, 0,
-           "closed tabs were implicitly cleared");
+  is(JSON.parse(ss.getClosedTabData(newWin)).length, 1,
+    "the closed tab was added before the window was overwritten");
+  delete newState.windows[0]._closedTabs;
+  delete newState.windows[0].sizemode;
 
-        is(newWin.windowState, newWin.STATE_MAXIMIZED,
-           "the window remains maximized");
-        newState.windows[0].sizemode = "normal";
-        ss.setWindowState(newWin, JSON.stringify(newState), true);
+  await setWindowState(newWin, newState, true);
+  await new Promise(resolve => newWin.setTimeout(resolve, 0));
 
-        newWin.setTimeout(function() {
-          isnot(newWin.windowState, newWin.STATE_MAXIMIZED,
-                "the window was explicitly unmaximized");
+  is(JSON.parse(ss.getClosedTabData(newWin)).length, 0,
+    "closed tabs were implicitly cleared");
 
-          BrowserTestUtils.closeWindow(newWin).then(finish);
-        }, 0);
-      }, 0);
-    }, 0);
-  });
-}
+  is(newWin.windowState, newWin.STATE_MAXIMIZED,
+    "the window remains maximized");
+  newState.windows[0].sizemode = "normal";
+
+  await setWindowState(newWin, newState, true);
+  await new Promise(resolve => newWin.setTimeout(resolve, 0));
+
+  isnot(newWin.windowState, newWin.STATE_MAXIMIZED,
+    "the window was explicitly unmaximized");
+
+  await BrowserTestUtils.closeWindow(newWin);
+});

browser/components/sessionstore/test/browser_490040.js

@@ -50,7 +50,7 @@ add_task(async function test_bug_490040() {
     let curClosedWindowCount = ss.getClosedWindowCount();
     let win = await BrowserTestUtils.openNewBrowserWindow();
 
-    ss.setWindowState(win, JSON.stringify(state.windowState), true);
+    await setWindowState(win, state.windowState, true);
     if (state.windowState.windows[0].tabs.length) {
       await BrowserTestUtils.browserLoaded(win.gBrowser.selectedBrowser);
     }

browser/components/sessionstore/test/browser_491577.js

@@ -2,12 +2,9 @@
  * 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/. */
 
-function test() {
+add_task(async function test_deleteClosedWindow() {
   /** Test for Bug 491577 **/
 
-  // test setup
-  waitForExplicitFinish();
-
   const REMEMBER = Date.now(), FORGET = Math.random();
   let test_state = {
     windows: [ { tabs: [{ entries: [{ url: "http://example.com/", triggeringPrincipal_base64 }] }], selected: 1 } ],
@@ -80,40 +77,39 @@ function test() {
 
   // open a window and add the above closed window list
   let newWin = openDialog(location, "_blank", "chrome,all,dialog=no");
-  promiseWindowLoaded(newWin).then(() => {
-    Services.prefs.setIntPref("browser.sessionstore.max_windows_undo",
-                              test_state._closedWindows.length);
-    ss.setWindowState(newWin, JSON.stringify(test_state), true);
+  await promiseWindowLoaded(newWin);
+  Services.prefs.setIntPref("browser.sessionstore.max_windows_undo",
+                            test_state._closedWindows.length);
+  await setWindowState(newWin, test_state, true);
 
-    let closedWindows = JSON.parse(ss.getClosedWindowData());
-    is(closedWindows.length, test_state._closedWindows.length,
-       "Closed window list has the expected length");
-    is(countByTitle(closedWindows, FORGET),
-       test_state._closedWindows.length - remember_count,
-       "The correct amount of windows are to be forgotten");
-    is(countByTitle(closedWindows, REMEMBER), remember_count,
-       "Everything is set up.");
+  let closedWindows = JSON.parse(ss.getClosedWindowData());
+  is(closedWindows.length, test_state._closedWindows.length,
+     "Closed window list has the expected length");
+  is(countByTitle(closedWindows, FORGET),
+     test_state._closedWindows.length - remember_count,
+     "The correct amount of windows are to be forgotten");
+  is(countByTitle(closedWindows, REMEMBER), remember_count,
+     "Everything is set up.");
 
-    // all of the following calls with illegal arguments should throw NS_ERROR_ILLEGAL_VALUE
-    ok(testForError(() => ss.forgetClosedWindow(-1)),
-       "Invalid window for forgetClosedWindow throws");
-    ok(testForError(() => ss.forgetClosedWindow(test_state._closedWindows.length + 1)),
-       "Invalid window for forgetClosedWindow throws");
+  // all of the following calls with illegal arguments should throw NS_ERROR_ILLEGAL_VALUE
+  ok(testForError(() => ss.forgetClosedWindow(-1)),
+     "Invalid window for forgetClosedWindow throws");
+  ok(testForError(() => ss.forgetClosedWindow(test_state._closedWindows.length + 1)),
+     "Invalid window for forgetClosedWindow throws");
 
-    // Remove third window, then first window
-    ss.forgetClosedWindow(2);
-    ss.forgetClosedWindow(null);
+  // Remove third window, then first window
+  ss.forgetClosedWindow(2);
+  ss.forgetClosedWindow(null);
 
-    closedWindows = JSON.parse(ss.getClosedWindowData());
-    is(closedWindows.length, remember_count,
-       "The correct amount of windows were removed");
-    is(countByTitle(closedWindows, FORGET), 0,
-       "All windows specifically forgotten were indeed removed");
-    is(countByTitle(closedWindows, REMEMBER), remember_count,
-       "... and windows not specifically forgetten weren't.");
+  closedWindows = JSON.parse(ss.getClosedWindowData());
+  is(closedWindows.length, remember_count,
+     "The correct amount of windows were removed");
+  is(countByTitle(closedWindows, FORGET), 0,
+     "All windows specifically forgotten were indeed removed");
+  is(countByTitle(closedWindows, REMEMBER), remember_count,
+     "... and windows not specifically forgetten weren't.");
 
-    // clean up
-    Services.prefs.clearUserPref("browser.sessionstore.max_windows_undo");
-    BrowserTestUtils.closeWindow(newWin).then(finish);
-  });
-}
+  // clean up
+  Services.prefs.clearUserPref("browser.sessionstore.max_windows_undo");
+  await BrowserTestUtils.closeWindow(newWin);
+});

browser/components/sessionstore/test/browser_495495.js

@@ -2,41 +2,33 @@
  * 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/. */
 
-function test() {
+add_task(async function test_urlbarFocus() {
   /** Test for Bug 495495 **/
 
-  waitForExplicitFinish();
-
   let newWin = openDialog(location, "_blank", "chrome,all,dialog=no,toolbar=yes");
-  promiseWindowLoaded(newWin).then(() => {
-    let state1 = ss.getWindowState(newWin);
-    BrowserTestUtils.closeWindow(newWin).then(() => {
+  await promiseWindowLoaded(newWin);
+  let state1 = ss.getWindowState(newWin);
+  await BrowserTestUtils.closeWindow(newWin);
 
-      newWin = openDialog(location, "_blank",
-        "chrome,extrachrome,menubar,resizable,scrollbars,status,toolbar=no,location,personal,directories,dialog=no");
-      promiseWindowLoaded(newWin).then(() => {
-        let state2 = ss.getWindowState(newWin);
+  newWin = openDialog(location, "_blank",
+    "chrome,extrachrome,menubar,resizable,scrollbars,status,toolbar=no,location,personal,directories,dialog=no");
+  await promiseWindowLoaded(newWin);
+  let state2 = ss.getWindowState(newWin);
 
-        function testState(state, expected, callback) {
-          let win = openDialog(location, "_blank", "chrome,all,dialog=no");
-          promiseWindowLoaded(win).then(() => {
+  async function testState(state, expected) {
+    let win = openDialog(location, "_blank", "chrome,all,dialog=no");
+    await promiseWindowLoaded(win);
 
-            is(win.gURLBar.readOnly, false,
-               "URL bar should not be read-only before setting the state");
-            ss.setWindowState(win, state, true);
-            is(win.gURLBar.readOnly, expected.readOnly,
-               "URL bar read-only state should be restored correctly");
+    is(win.gURLBar.readOnly, false,
+       "URL bar should not be read-only before setting the state");
+    await setWindowState(win, state, true);
+    is(win.gURLBar.readOnly, expected.readOnly,
+       "URL bar read-only state should be restored correctly");
 
-            BrowserTestUtils.closeWindow(win).then(callback);
-          });
-        }
+    await BrowserTestUtils.closeWindow(win);
+  }
 
-        BrowserTestUtils.closeWindow(newWin).then(() => {
-          testState(state1, {readOnly: false}, function() {
-            testState(state2, {readOnly: true}, finish);
-          });
-        });
-      });
-    });
-  });
-}
+  await BrowserTestUtils.closeWindow(newWin);
+  await testState(state1, {readOnly: false});
+  await testState(state2, {readOnly: true});
+});

browser/components/sessionstore/test/browser_514751.js

@@ -2,11 +2,9 @@
  * 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/. */
 
-function test() {
+add_task(async function test_malformedURI() {
   /** Test for Bug 514751 (Wallpaper) **/
 
-  waitForExplicitFinish();
-
   let state = {
     windows: [{
       tabs: [{
@@ -19,18 +17,16 @@ function test() {
   };
 
   var theWin = openDialog(location, "", "chrome,all,dialog=no");
-  theWin.addEventListener("load", function() {
-    executeSoon(function() {
-      var gotError = false;
-      try {
-        ss.setWindowState(theWin, JSON.stringify(state), true);
-      } catch (e) {
-        if (/NS_ERROR_MALFORMED_URI/.test(e))
-          gotError = true;
-      }
-      ok(!gotError, "Didn't get a malformed URI error.");
-      BrowserTestUtils.closeWindow(theWin).then(finish);
-    });
-  }, {once: true});
-}
+  await promiseWindowLoaded(theWin);
 
+  var gotError = false;
+  try {
+    await setWindowState(theWin, state, true);
+  } catch (e) {
+    if (/NS_ERROR_MALFORMED_URI/.test(e))
+      gotError = true;
+  }
+
+  ok(!gotError, "Didn't get a malformed URI error.");
+  await BrowserTestUtils.closeWindow(theWin);
+});

browser/components/sessionstore/test/browser_607016.js

@@ -82,7 +82,7 @@ add_task(async function() {
   }
 
   // Set the test state.
-  ss.setBrowserState(JSON.stringify(state));
+  await setBrowserState(state);
 
   // Wait until the selected tab is restored and all others are pending.
   await Promise.all(Array.map(gBrowser.tabs, tab => {

browser/components/sessionstore/test/browser_637020.js

@@ -42,6 +42,7 @@ add_task(async function test() {
   let backupState = SessionStore.getBrowserState();
   SessionStore.setBrowserState(JSON.stringify(TEST_STATE));
   let win = await promiseWindow;
+  await promiseWindowRestored(win);
 
   // The window has now been opened. Check the state that is returned,
   // this should come from the cache while the window isn't restored, yet.

browser/components/sessionstore/test/browser_687710.js

@@ -37,12 +37,12 @@ var state = {windows: [{tabs: [{entries: [
   }
 ]}]}]};
 
-function test() {
+add_task(async function test() {
   registerCleanupFunction(function() {
     ss.setBrowserState(stateBackup);
   });
 
   /* This test fails by hanging. */
-  ss.setBrowserState(JSON.stringify(state));
+  await setBrowserState(state);
   ok(true, "Didn't hang!");
-}
+});

browser/components/sessionstore/test/browser_694378.js

@@ -5,20 +5,17 @@
 // 1.  call ss.setWindowState with a broken state
 // 1a. ensure that it doesn't throw.
 
-function test() {
-  waitForExplicitFinish();
-
+add_task(async function test_brokenWindowState() {
   let brokenState = {
     windows: [
       { tabs: [{ entries: [{ url: "about:mozilla", triggeringPrincipal_base64 }] }] }
     ],
     selectedWindow: 2
   };
-  let brokenStateString = JSON.stringify(brokenState);
 
   let gotError = false;
   try {
-    ss.setWindowState(window, brokenStateString, true);
+    await setWindowState(window, brokenState, true);
   } catch (ex) {
     gotError = true;
     info(ex);
@@ -28,5 +25,5 @@ function test() {
 
   // Make sure that we reset the state. Use a full state just in case things get crazy.
   let blankState = { windows: [{ tabs: [{ entries: [{ url: "about:blank", triggeringPrincipal_base64 }] }]}]};
-  waitForBrowserState(blankState, finish);
-}
+  await promiseBrowserState(blankState);
+});

browser/components/sessionstore/test/browser_frame_history.js

@@ -101,14 +101,14 @@ add_task(async function() {
 // Now, test that we don't record history if the iframe is added dynamically
 add_task(async function() {
   // Start with an empty history
-    let blankState = JSON.stringify({
-      windows: [{
-        tabs: [{ entries: [{ url: "about:blank", triggeringPrincipal_base64 }] }],
-        _closedTabs: []
-      }],
-      _closedWindows: []
-    });
-    ss.setBrowserState(blankState);
+  let blankState = JSON.stringify({
+    windows: [{
+      tabs: [{ entries: [{ url: "about:blank", triggeringPrincipal_base64 }] }],
+      _closedTabs: []
+    }],
+    _closedWindows: []
+  });
+  await setBrowserState(blankState);
 
   let testURL = getRootDirectory(gTestPath) + "browser_frame_history_index_blank.html";
   let tab = BrowserTestUtils.addTab(gBrowser, testURL);

browser/components/sessionstore/test/browser_merge_closed_tabs.js

@@ -36,12 +36,12 @@ add_task(async function() {
 
   // Open a new window and restore it to an initial state.
   let win = await promiseNewWindowLoaded({private: false});
-  SessionStore.setWindowState(win, JSON.stringify(initialState), true);
+  await setWindowState(win, initialState, true);
   is(SessionStore.getClosedTabCount(win), 2, "2 closed tabs after restoring initial state");
 
   // Restore the new state but do not overwrite existing tabs (this should
   // cause the closed tabs to be merged).
-  SessionStore.setWindowState(win, JSON.stringify(restoreState), false);
+  await setWindowState(win, restoreState);
 
   // Verify the windows closed tab data is correct.
   let iClosed = initialState.windows[0]._closedTabs;

browser/components/sessionstore/test/browser_remoteness_flip_on_restore.js

@@ -134,7 +134,7 @@ async function runScenarios(scenarios) {
     let state = prepareState(scenario.stateToRestore,
                              scenario.selectedTab);
 
-    SessionStore.setWindowState(win, state, true);
+    await setWindowState(win, state, true);
 
     for (let i = 0; i < scenario.expectedRemoteness.length; ++i) {
       let expectedRemoteness = scenario.expectedRemoteness[i];

browser/components/sessionstore/test/browser_restore_cookies_noOriginAttributes.js

@@ -134,7 +134,7 @@ add_task(async function run_test() {
   let win = await promiseNewWindowLoaded();
 
   // Restore window with session cookies that have no originAttributes.
-  ss.setWindowState(win, SESSION_DATA, true);
+  await setWindowState(win, SESSION_DATA, true);
 
   let enumerator = Services.cookies.getCookiesFromHost(TEST_HOST, {});
   let cookie;
@@ -154,7 +154,7 @@ add_task(async function run_test() {
   Services.cookies.removeAll();
 
   // Restore window with session cookies that have originAttributes within.
-  ss.setWindowState(win, SESSION_DATA_OA, true);
+  await setWindowState(win, SESSION_DATA_OA, true);
 
   enumerator = Services.cookies.getCookiesFromHost(TEST_HOST, {});
   cookieCount = 0;

browser/components/sessionstore/test/browser_speculative_connect.js

@@ -29,7 +29,7 @@ add_task(async function speculative_connect_restore_on_demand() {
   // Reopen a window.
   let newWin = undoCloseWindow(0);
   // Make sure we wait until this window is restored.
-  await BrowserTestUtils.waitForEvent(newWin, "load");
+  await promiseWindowRestored(newWin);
 
   let tabs = newWin.gBrowser.tabs;
   is(tabs.length, TEST_URLS.length + 1, "Restored right number of tabs");
@@ -79,7 +79,7 @@ add_task(async function speculative_connect_restore_automatically() {
   // Reopen a window.
   let newWin = undoCloseWindow(0);
   // Make sure we wait until this window is restored.
-  await BrowserTestUtils.waitForEvent(newWin, "load");
+  await promiseWindowRestored(newWin);
 
   let tabs = newWin.gBrowser.tabs;
   is(tabs.length, TEST_URLS.length + 1, "Restored right number of tabs");

browser/components/sessionstore/test/browser_windowStateContainer.js

@@ -41,7 +41,7 @@ add_task(async function() {
     await TabStateFlusher.flush(tab.linkedBrowser);
   }
 
-  ss.setWindowState(win2, JSON.stringify(winState), true);
+  await setWindowState(win2, winState, true);
 
   for (let i = 0; i < 4; i++) {
     let browser = win2.gBrowser.tabs[i].linkedBrowser;
@@ -101,7 +101,7 @@ add_task(async function() {
   win2.gBrowser.moveTabTo(win2.gBrowser.tabs[0], win2.gBrowser.tabs.length - 1);
   await TabStateFlusher.flush(win2.gBrowser.tabs[0].linkedBrowser);
 
-  ss.setWindowState(win2, JSON.stringify(winState), true);
+  await setWindowState(win2, winState, true);
 
   for (let i = 0; i < 2; i++) {
     let browser = win2.gBrowser.tabs[i].linkedBrowser;

browser/components/sessionstore/test/head.js

@@ -181,6 +181,20 @@ function promiseTabState(tab, state) {
   return promise;
 }
 
+function promiseWindowRestored(win) {
+  return new Promise(resolve => win.addEventListener("SSWindowRestored", resolve, {once: true}));
+}
+
+async function setBrowserState(state, win = window) {
+  ss.setBrowserState(typeof state != "string" ? JSON.stringify(state) : state);
+  await promiseWindowRestored(win);
+}
+
+async function setWindowState(win, state, overwrite = false) {
+  ss.setWindowState(win, typeof state != "string" ? JSON.stringify(state) : state, overwrite);
+  await promiseWindowRestored(win);
+}
+
 /**
  * Wait for a content -> chrome message.
  */