From c64d63e759cf59aad8f74af56bfe0b7d472fa9ae Mon Sep 17 00:00:00 2001
From: Justin Ridgewell <jridgewell@google.com>
Date: Wed, 26 Apr 2017 18:05:16 -0400
Subject: [PATCH] Simplfy expectMessageFromParent callback

---
 testing/iframe.js | 27 +++++++++++++--------------
 1 file changed, 13 insertions(+), 14 deletions(-)

diff --git a/testing/iframe.js b/testing/iframe.js
index c74e37e5d7ca6..183b3b0d1f8ee 100644
--- a/testing/iframe.js
+++ b/testing/iframe.js
@@ -319,13 +319,19 @@ export function createIframeWithMessageStub(win) {
   /**
    * Returns a Promise that resolves when the iframe acknowledged the reception
    * of the specified message.
+   * @param {function(?Object, !Object|string)|string} callbackOrType
+   *     A callback that determines if this is the message we expected. If a
+   *     string is passed, the determination is based on whether the message's
+   *     type matches the string.
    */
-  element.expectMessageFromParent = (type, callback) => {
-    if (typeof type === 'function') {
-      callback = type;
-      type = '';
-    } else if (!callback) {
-      callback = returnTrue;
+  element.expectMessageFromParent = (callbackOrType) => {
+    let filter;
+    if (typeof callbackOrType === 'string') {
+      filter = (data) => {
+        return 'type' in data && data.type == type;
+      };
+    } else {
+      filter = callbackOrType;
     }
 
     return new Promise((resolve, reject) => {
@@ -335,11 +341,8 @@ export function createIframeWithMessageStub(win) {
         }
         const message = event.data.receivedMessage;
         const data = parseIfNeeded(message);
-        if (type && 'type' in data && data.type !== type) {
-          return;
-        }
         try {
-          if (callback(data, message)) {
+          if (filter(data, message)) {
             win.removeEventListener('message', listener);
             resolve(data || message);
           }
@@ -352,10 +355,6 @@ export function createIframeWithMessageStub(win) {
     });
   };
 
-  function returnTrue() {
-    return true;
-  }
-
   return element;
 }