Bug 1504756 - [marionette] Use waitForMessage() to wait for an expected message manager message. r=ato

Depends on D13660

Differential Revision: https://phabricator.services.mozilla.com/D13661

--HG--
extra : moz-landing-system : lando

testing/marionette/doc/internals/sync.rst

@@ -18,3 +18,5 @@ Provides an assortment of synchronisation primitives.
   :members:
 
 .. js:autofunction:: waitForEvent
+
+.. js:autofunction:: waitForMessage

testing/marionette/sync.js

@@ -13,6 +13,7 @@ const {
   stack,
   TimeoutError,
 } = ChromeUtils.import("chrome://marionette/content/error.js", {});
+const {truncate} = ChromeUtils.import("chrome://marionette/content/format.js", {});
 const {Log} = ChromeUtils.import("chrome://marionette/content/log.js", {});
 
 XPCOMUtils.defineLazyGetter(this, "log", Log.get);
@@ -26,6 +27,7 @@ this.EXPORTED_SYMBOLS = [
   "Sleep",
   "TimedPromise",
   "waitForEvent",
+  "waitForMessage",
 ];
 
 const {TYPE_ONE_SHOT, TYPE_REPEATING_SLACK} = Ci.nsITimer;
@@ -464,3 +466,48 @@ function waitForEvent(subject, eventName,
     }, capture, wantsUntrusted);
   });
 }
+
+/**
+ * Wait for a message to be fired from a particular message manager.
+ *
+ * This method has been duplicated from BrowserTestUtils.jsm.
+ *
+ * @param {nsIMessageManager} messageManager
+ *     The message manager that should be used.
+ * @param {string} messageName
+ *     The message to wait for.
+ * @param {Object=} options
+ *     Extra options.
+ * @param {function(Message)=} options.checkFn
+ *     Called with the ``Message`` object as argument, should return ``true``
+ *     if the message is the expected one, or ``false`` if it should be
+ *     ignored and listening should continue. If not specified, the first
+ *     message with the specified name resolves the returned promise.
+ *
+ * @return {Promise.<Object>}
+ *     Promise which resolves to the data property of the received
+ *     ``Message``.
+ */
+function waitForMessage(messageManager, messageName,
+    {checkFn = undefined} = {}) {
+  if (messageManager == null || !("addMessageListener" in messageManager)) {
+    throw new TypeError();
+  }
+  if (typeof messageName != "string") {
+    throw new TypeError();
+  }
+  if (checkFn && typeof checkFn != "function") {
+    throw new TypeError();
+  }
+
+  return new Promise(resolve => {
+    messageManager.addMessageListener(messageName, function onMessage(msg) {
+      log.trace(`Received ${messageName} for ${msg.target}`);
+      if (checkFn && !checkFn(msg)) {
+        return;
+      }
+      messageManager.removeMessageListener(messageName, onMessage);
+      resolve(msg.data);
+    });
+  });
+}

testing/marionette/test/unit/test_sync.js

@@ -9,6 +9,7 @@ const {
   Sleep,
   TimedPromise,
   waitForEvent,
+  waitForMessage,
 } = ChromeUtils.import("chrome://marionette/content/sync.js", {});
 
 const DEFAULT_TIMEOUT = 2000;
@@ -55,6 +56,36 @@ class MockElement {
   }
 }
 
+/**
+ * Mimic a message manager for sending messages.
+ */
+class MessageManager {
+  constructor() {
+    this.func = null;
+    this.message = null;
+  }
+
+  addMessageListener(message, func) {
+    this.func = func;
+    this.message = message;
+  }
+
+  removeMessageListener(message) {
+    this.func = null;
+    this.message = null;
+  }
+
+  send(message, data) {
+    if (this.func) {
+      this.func({
+        data,
+        message,
+        target: this,
+      });
+    }
+  }
+}
+
 /**
  * Mimics nsITimer, but instead of using a system clock you can
  * preprogram it to invoke the callback after a given number of ticks.
@@ -353,3 +384,42 @@ add_task(async function test_waitForEvent_wantsUntrustedTypes() {
     equal(expected_untrusted, event.untrusted);
   }
 });
+
+add_task(async function test_waitForMessage_messageManagerAndMessageTypes() {
+  let messageManager = new MessageManager();
+
+  for (let manager of ["foo", 42, null, undefined, true, [], {}]) {
+    Assert.throws(() => waitForMessage(manager, "message"), /TypeError/);
+  }
+
+  for (let message of [42, null, undefined, true, [], {}]) {
+    Assert.throws(() => waitForEvent(messageManager, message), /TypeError/);
+  }
+
+  let data = {"foo": "bar"};
+  let sent = waitForMessage(messageManager, "message");
+  messageManager.send("message", data);
+  equal(data, await sent);
+});
+
+add_task(async function test_waitForMessage_checkFnTypes() {
+  let messageManager = new MessageManager();
+
+  for (let checkFn of ["foo", 42, true, [], {}]) {
+    Assert.throws(() => waitForMessage(
+        messageManager, "message", {checkFn}), /TypeError/);
+  }
+
+  let data1 = {"fo": "bar"};
+  let data2 = {"foo": "bar"};
+
+  for (let checkFn of [null, undefined, msg => "foo" in msg.data]) {
+    let expected_data = (checkFn == null) ? data1 : data2;
+
+    messageManager = new MessageManager();
+    let sent = waitForMessage(messageManager, "message", {checkFn});
+    messageManager.send("message", data1);
+    messageManager.send("message", data2);
+    equal(expected_data, await sent);
+  }
+});