From 1a9ec135572421f7ba565db2422fe84dbb145ff7 Mon Sep 17 00:00:00 2001
From: Simon Giesecke <sgiesecke@mozilla.com>
Date: Wed, 30 Oct 2019 15:26:25 +0000
Subject: [PATCH] Bug 1586189 - Re-implement waitForMessage based on
 waitForContentEvent; r=mccr8,janv

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

--HG--
extra : moz-landing-system : lando
---
 dom/indexedDB/test/head.js | 53 +++++++++++++++-----------------------
 dom/quota/test/head.js     | 51 +++++++++++++++---------------------
 2 files changed, 42 insertions(+), 62 deletions(-)

diff --git a/dom/indexedDB/test/head.js b/dom/indexedDB/test/head.js
index a7251a5298e9..657fbcd9a51a 100644
--- a/dom/indexedDB/test/head.js
+++ b/dom/indexedDB/test/head.js
@@ -79,40 +79,29 @@ function dismissNotification(popup) {
 }
 
 function waitForMessage(aMessage, browser) {
-  return new Promise((resolve, reject) => {
-    /* eslint-disable no-undef */
-    // When contentScript runs, "this" is a ContentFrameMessageManager (so that's where
-    // addEventListener will add the listener), but the non-bubbling "message" event is
-    // sent to the Window involved, so we need a capturing listener.
-    function contentScript() {
-      addEventListener(
-        "message",
-        function(event) {
-          sendAsyncMessage("testLocal:message", { message: event.data });
-        },
-        { once: true, capture: true },
-        true
-      );
-    }
-    /* eslint-enable no-undef */
-
-    let script = "data:,(" + contentScript.toString() + ")();";
-
-    let mm = browser.selectedBrowser.messageManager;
-
-    mm.addMessageListener("testLocal:message", function listener(msg) {
-      mm.removeMessageListener("testLocal:message", listener);
-      mm.removeDelayedFrameScript(script);
-      is(msg.data.message, aMessage, "received " + aMessage);
-      if (msg.data.message == aMessage) {
-        resolve();
-      } else {
-        reject();
+  // We cannot capture aMessage inside the checkFn, so we override the
+  // checkFn.toSource to tunnel aMessage instead.
+  let checkFn = function() {};
+  checkFn.toSource = function() {
+    return `function checkFn(event) {
+      let message = ${aMessage.toSource()};
+      is(event.data, message, "Received: " + message);
+      if (event.data == message) {
+        return true;
       }
-    });
+      throw new Error(
+       \`Unexpected result: \$\{event.data\}, expected \$\{message\}\`
+      );
+    }`;
+  };
 
-    mm.loadFrameScript(script, true);
-  });
+  return BrowserTestUtils.waitForContentEvent(
+    browser.selectedBrowser,
+    "message",
+    /* capture */ true,
+    checkFn,
+    /* wantsUntrusted */ true
+  );
 }
 
 function dispatchEvent(eventName) {
diff --git a/dom/quota/test/head.js b/dom/quota/test/head.js
index 395d93b5b21c..b8da5f7bdea2 100644
--- a/dom/quota/test/head.js
+++ b/dom/quota/test/head.js
@@ -101,38 +101,29 @@ function dismissNotification(popup, win) {
 }
 
 function waitForMessage(aMessage, browser) {
-  return new Promise((resolve, reject) => {
-    // When contentScript runs, "this" is a ContentFrameMessageManager (so that's where
-    // addEventListener will add the listener), but the non-bubbling "message" event is
-    // sent to the Window involved, so we need a capturing listener.
-    function contentScript() {
-      addEventListener(
-        "message",
-        function(event) {
-          sendAsyncMessage("testLocal:persisted", { persisted: event.data });
-        },
-        { once: true, capture: true },
-        true
-      );
-    }
-
-    let script = "data:,(" + contentScript.toString() + ")();";
-
-    let mm = browser.selectedBrowser.messageManager;
-
-    mm.addMessageListener("testLocal:persisted", function listener(msg) {
-      mm.removeMessageListener("testLocal:persisted", listener);
-      mm.removeDelayedFrameScript(script);
-      is(msg.data.persisted, aMessage, "received " + aMessage);
-      if (msg.data.persisted == aMessage) {
-        resolve();
-      } else {
-        reject();
+  // We cannot capture aMessage inside the checkFn, so we override the
+  // checkFn.toSource to tunnel aMessage instead.
+  let checkFn = function() {};
+  checkFn.toSource = function() {
+    return `function checkFn(event) {
+      let message = ${aMessage.toSource()};
+      is(event.data, message, "Received: " + message);
+      if (event.data == message) {
+        return true;
       }
-    });
+      throw new Error(
+       \`Unexpected result: \$\{event.data\}, expected \$\{message\}\`
+      );
+    }`;
+  };
 
-    mm.loadFrameScript(script, true);
-  });
+  return BrowserTestUtils.waitForContentEvent(
+    browser.selectedBrowser,
+    "message",
+    /* capture */ true,
+    checkFn,
+    /* wantsUntrusted */ true
+  );
 }
 
 function removePermission(url, permission) {