From 7389b135f9ab93ef843f0695d1ed834aa16bce6f Mon Sep 17 00:00:00 2001
From: Tim Taubert <ttaubert@mozilla.com>
Date: Tue, 9 Jun 2015 22:11:43 +0200
Subject: [PATCH] Bug 1167508 - Fix BrowserTestUtils.browserLoaded() and use it
 in sessionstore/head.js r=billm

---
 browser/components/sessionstore/test/head.js            | 9 +--------
 testing/mochitest/BrowserTestUtils/BrowserTestUtils.jsm | 9 ++++-----
 2 files changed, 5 insertions(+), 13 deletions(-)

diff --git a/browser/components/sessionstore/test/head.js b/browser/components/sessionstore/test/head.js
index 7182e591f99e..9519cc272449 100644
--- a/browser/components/sessionstore/test/head.js
+++ b/browser/components/sessionstore/test/head.js
@@ -282,14 +282,7 @@ let promiseForEachSessionRestoreFile = Task.async(function*(cb) {
 });
 
 function promiseBrowserLoaded(aBrowser, ignoreSubFrames = true) {
-  return new Promise(resolve => {
-    aBrowser.messageManager.addMessageListener("ss-test:loadEvent", function onLoad(msg) {
-      if (!ignoreSubFrames || !msg.data.subframe) {
-        aBrowser.messageManager.removeMessageListener("ss-test:loadEvent", onLoad);
-        resolve();
-      }
-    });
-  });
+  return BrowserTestUtils.browserLoaded(aBrowser, !ignoreSubFrames);
 }
 
 function whenWindowLoaded(aWindow, aCallback = next) {
diff --git a/testing/mochitest/BrowserTestUtils/BrowserTestUtils.jsm b/testing/mochitest/BrowserTestUtils/BrowserTestUtils.jsm
index ed753d36df34..21dd47030681 100644
--- a/testing/mochitest/BrowserTestUtils/BrowserTestUtils.jsm
+++ b/testing/mochitest/BrowserTestUtils/BrowserTestUtils.jsm
@@ -140,11 +140,10 @@ this.BrowserTestUtils = {
    */
   browserLoaded(browser, includeSubFrames=false) {
     return new Promise(resolve => {
-      browser.messageManager.addMessageListener("browser-test-utils:loadEvent",
-                                                 function onLoad(msg) {
-        if (!msg.data.subframe || includeSubFrames) {
-          browser.messageManager.removeMessageListener(
-            "browser-test-utils:loadEvent", onLoad);
+      let mm = browser.ownerDocument.defaultView.messageManager;
+      mm.addMessageListener("browser-test-utils:loadEvent", function onLoad(msg) {
+        if (msg.target == browser && (!msg.data.subframe || includeSubFrames)) {
+          mm.removeMessageListener("browser-test-utils:loadEvent", onLoad);
           resolve();
         }
       });