Bug 1412456 - Add BrowserTestUtils.addContentEventListener (r=mconley)

This function makes it possible to listen for multiple events from the
content process, even when there are frameloader swaps.

This commit also adds a checkFn param to firstBrowserLoaded, which is
useful.

MozReview-Commit-ID: 93ItHIPSGVU

testing/mochitest/BrowserTestUtils/BrowserTestUtils.jsm

@@ -38,6 +38,8 @@ const EXISTING_JSID = Cc[PROCESSSELECTOR_CONTRACTID];
 const DEFAULT_PROCESSSELECTOR_CID = EXISTING_JSID ?
   Components.ID(EXISTING_JSID.number) : null;
 
+let gListenerId = 0;
+
 // A process selector that always asks for a new process.
 function NewProcessSelector() {
 }
@@ -301,16 +303,26 @@ this.BrowserTestUtils = {
    * loaded its DOM yet, and where you can't easily use browserLoaded
    * on gBrowser.selectedBrowser since gBrowser doesn't yet exist.
    *
-   * @param {win}
+   * @param {xul:window} window
    *        A newly opened window for which we're waiting for the
    *        first browser load.
+   * @param {Boolean} aboutBlank [optional]
+   *        If false, about:blank loads are ignored and we continue
+   *        to wait.
+   * @param {function or null} checkFn [optional]
+   *        If checkFn(browser) returns false, the load is ignored
+   *        and we continue to wait.
    *
    * @return {Promise}
    * @resolves Once the selected browser fires its load event.
    */
-  firstBrowserLoaded(win, aboutBlank = true) {
+  firstBrowserLoaded(win, aboutBlank = true, checkFn = null) {
     let mm = win.messageManager;
     return this.waitForMessage(mm, "browser-test-utils:loadEvent", (msg) => {
+      if (checkFn) {
+        return checkFn(msg.target);
+      }
+
       let selectedBrowser = win.gBrowser.selectedBrowser;
       return msg.target == selectedBrowser &&
              (aboutBlank || selectedBrowser.currentURI.spec != "about:blank")
@@ -826,6 +838,115 @@ this.BrowserTestUtils = {
         });
   },
 
+  /**
+   * Adds a content event listener on the given browser
+   * element. Similar to waitForContentEvent, but the listener will
+   * fire until it is removed. A callable object is returned that,
+   * when called, removes the event listener. Note that this function
+   * works even if the browser's frameloader is swapped.
+   *
+   * @param {xul:browser} browser
+   *        The browser element to listen for events in.
+   * @param {string} eventName
+   *        Name of the event to listen to.
+   * @param {function} listener
+   *        Function to call in parent process when event fires.
+   *        Not passed any arguments.
+   * @param {bool} useCapture [optional]
+   *        Whether to use a capturing listener.
+   * @param {function} checkFn [optional]
+   *        Called with the Event object as argument, should return true if the
+   *        event is the expected one, or false if it should be ignored and
+   *        listening should continue. If not specified, the first event with
+   *        the specified name resolves the returned promise. This is called
+   *        within the content process and can have no closure environment.
+   * @param {bool} wantsUntrusted [optional]
+   *        Whether to accept untrusted events
+   * @param {bool} autoremove [optional]
+   *        Whether the listener should be removed when |browser| is removed
+   *        from the DOM. Note that, if this flag is true, it won't be possible
+   *        to listen for events after a frameloader swap.
+   *
+   * @returns function
+   *        If called, the return value will remove the event listener.
+   */
+  addContentEventListener(browser,
+                          eventName,
+                          listener,
+                          useCapture = false,
+                          checkFn,
+                          wantsUntrusted = false,
+                          autoremove = true) {
+    let id = gListenerId++;
+    let checkFnSource = checkFn ? encodeURIComponent(escape(checkFn.toSource())) : "";
+
+    // To correctly handle frameloader swaps, we load a frame script
+    // into all tabs but ignore messages from the ones not related to
+    // |browser|.
+
+    function frameScript(id, eventName, useCapture, checkFnSource, wantsUntrusted) {
+      let checkFn;
+      if (checkFnSource) {
+        checkFn = eval(`(() => (${unescape(checkFnSource)}))()`);
+      }
+
+      function listener(event) {
+        if (checkFn && !checkFn(event)) {
+          return;
+        }
+        sendAsyncMessage("ContentEventListener:Run", id);
+      }
+      function removeListener(msg) {
+        if (msg.data == id) {
+          removeMessageListener("ContentEventListener:Remove", removeListener);
+          removeEventListener(eventName, listener, useCapture, wantsUntrusted);
+        }
+      }
+      addMessageListener("ContentEventListener:Remove", removeListener);
+      addEventListener(eventName, listener, useCapture, wantsUntrusted);
+    }
+
+    let frameScriptSource =
+        `data:,(${frameScript.toString()})(${id}, "${eventName}", ${useCapture}, "${checkFnSource}", ${wantsUntrusted})`;
+
+    let mm = Services.mm;
+
+    function runListener(msg) {
+      if (msg.data == id && msg.target == browser) {
+        listener();
+      }
+    }
+    mm.addMessageListener("ContentEventListener:Run", runListener);
+
+    let needCleanup = true;
+
+    let unregisterFunction = function() {
+      if (!needCleanup) {
+        return;
+      }
+      needCleanup = false;
+      mm.removeMessageListener("ContentEventListener:Run", runListener);
+      mm.broadcastAsyncMessage("ContentEventListener:Remove", id);
+      mm.removeDelayedFrameScript(frameScriptSource);
+      if (autoremove) {
+        Services.obs.removeObserver(cleanupObserver, "message-manager-close");
+      }
+    };
+
+    function cleanupObserver(subject, topic, data) {
+      if (subject == browser.messageManager) {
+        unregisterFunction();
+      }
+    }
+    if (autoremove) {
+      Services.obs.addObserver(cleanupObserver, "message-manager-close");
+    }
+
+    mm.loadFrameScript(frameScriptSource, true);
+
+    return unregisterFunction;
+  },
+
   /**
    * Like browserLoaded, but waits for an error page to appear.
    * This explicitly deals with cases where the browser is not currently remote and a