diff --git a/testing/marionette/doc/internals/sync.rst b/testing/marionette/doc/internals/sync.rst index ae018afea2cf..d97c12742781 100644 --- a/testing/marionette/doc/internals/sync.rst +++ b/testing/marionette/doc/internals/sync.rst @@ -18,3 +18,5 @@ Provides an assortment of synchronisation primitives. :members: .. js:autofunction:: waitForEvent + +.. js:autofunction:: waitForMessage diff --git a/testing/marionette/sync.js b/testing/marionette/sync.js index badfea82667e..d88ed8ac02cc 100644 --- a/testing/marionette/sync.js +++ b/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.} + * 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); + }); + }); +} diff --git a/testing/marionette/test/unit/test_sync.js b/testing/marionette/test/unit/test_sync.js index 795dc537b964..aaae9b9302ec 100644 --- a/testing/marionette/test/unit/test_sync.js +++ b/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); + } +});