From caaf4d2086036d4bf201587529dd3e3b38b6f9c6 Mon Sep 17 00:00:00 2001
From: Islam ElHakmi <eslam.elhakmey3@gmail.com>
Date: Tue, 19 May 2020 09:09:31 +0200
Subject: [PATCH] fix: support async functions as an argument for
 waitForFunction (#5682)

---
 docs/api.md           | 16 ++++++++++++++++
 src/DOMWorld.ts       | 24 +++++++++++------------
 test/waittask.spec.js | 44 +++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 72 insertions(+), 12 deletions(-)

diff --git a/docs/api.md b/docs/api.md
index 686dc2f6..bc4ad17e 100644
--- a/docs/api.md
+++ b/docs/api.md
@@ -2098,6 +2098,22 @@ const selector = '.foo';
 await page.waitForFunction(selector => !!document.querySelector(selector), {}, selector);
 ```
 
+The predicate of `page.waitForFunction` can be asynchronous too:
+
+```js
+const username = 'github-username';
+await page.waitForFunction(async username => {
+  const githubResponse = await fetch(`https://api.github.com/users/${username}`);
+  const githubUser = await githubResponse.json();
+  // show the avatar
+  const img = document.createElement('img');
+  img.src = githubUser.avatar_url;
+  // wait 3 seconds
+  await new Promise((resolve, reject) => setTimeout(resolve, 3000));
+  img.remove();
+}, {}, username);
+```
+
 Shortcut for [page.mainFrame().waitForFunction(pageFunction[, options[, ...args]])](#framewaitforfunctionpagefunction-options-args).
 
 #### page.waitForNavigation([options])
diff --git a/src/DOMWorld.ts b/src/DOMWorld.ts
index 43ad349d..886b2b49 100644
--- a/src/DOMWorld.ts
+++ b/src/DOMWorld.ts
@@ -692,18 +692,18 @@ async function waitForPredicatePageFunction(
   /**
    * @return {!Promise<*>}
    */
-  function pollMutation(): Promise<unknown> {
-    const success = predicate(...args);
+  async function pollMutation(): Promise<unknown> {
+    const success = await predicate(...args);
     if (success) return Promise.resolve(success);
 
     let fulfill;
     const result = new Promise((x) => (fulfill = x));
-    const observer = new MutationObserver(() => {
+    const observer = new MutationObserver(async () => {
       if (timedOut) {
         observer.disconnect();
         fulfill();
       }
-      const success = predicate(...args);
+      const success = await predicate(...args);
       if (success) {
         observer.disconnect();
         fulfill(success);
@@ -717,35 +717,35 @@ async function waitForPredicatePageFunction(
     return result;
   }
 
-  function pollRaf(): Promise<unknown> {
+  async function pollRaf(): Promise<unknown> {
     let fulfill;
     const result = new Promise((x) => (fulfill = x));
-    onRaf();
+    await onRaf();
     return result;
 
-    function onRaf(): void {
+    async function onRaf(): Promise<unknown> {
       if (timedOut) {
         fulfill();
         return;
       }
-      const success = predicate(...args);
+      const success = await predicate(...args);
       if (success) fulfill(success);
       else requestAnimationFrame(onRaf);
     }
   }
 
-  function pollInterval(pollInterval: number): Promise<unknown> {
+  async function pollInterval(pollInterval: number): Promise<unknown> {
     let fulfill;
     const result = new Promise((x) => (fulfill = x));
-    onTimeout();
+    await onTimeout();
     return result;
 
-    function onTimeout(): void {
+    async function onTimeout(): Promise<unknown> {
       if (timedOut) {
         fulfill();
         return;
       }
-      const success = predicate(...args);
+      const success = await predicate(...args);
       if (success) fulfill(success);
       else setTimeout(onTimeout, pollInterval);
     }
diff --git a/test/waittask.spec.js b/test/waittask.spec.js
index a0c3b52e..1d604a65 100644
--- a/test/waittask.spec.js
+++ b/test/waittask.spec.js
@@ -138,6 +138,22 @@ describe('waittask specs', function () {
       await watchdog;
       expect(Date.now() - startTime).not.toBeLessThan(polling / 2);
     });
+    it('should poll on interval async', async () => {
+      const { page } = getTestState();
+      let success = false;
+      const startTime = Date.now();
+      const polling = 100;
+      const watchdog = page
+        .waitForFunction(async () => window.__FOO === 'hit', { polling })
+        .then(() => (success = true));
+      await page.evaluate(async () => (window.__FOO = 'hit'));
+      expect(success).toBe(false);
+      await page.evaluate(async () =>
+        document.body.appendChild(document.createElement('div'))
+      );
+      await watchdog;
+      expect(Date.now() - startTime).not.toBeLessThan(polling / 2);
+    });
     it('should poll on mutation', async () => {
       const { page } = getTestState();
 
@@ -152,6 +168,22 @@ describe('waittask specs', function () {
       );
       await watchdog;
     });
+    it('should poll on mutation async', async () => {
+      const { page } = getTestState();
+
+      let success = false;
+      const watchdog = page
+        .waitForFunction(async () => window.__FOO === 'hit', {
+          polling: 'mutation',
+        })
+        .then(() => (success = true));
+      await page.evaluate(async () => (window.__FOO = 'hit'));
+      expect(success).toBe(false);
+      await page.evaluate(async () =>
+        document.body.appendChild(document.createElement('div'))
+      );
+      await watchdog;
+    });
     it('should poll on raf', async () => {
       const { page } = getTestState();
 
@@ -161,6 +193,18 @@ describe('waittask specs', function () {
       await page.evaluate(() => (window.__FOO = 'hit'));
       await watchdog;
     });
+    it('should poll on raf async', async () => {
+      const { page } = getTestState();
+
+      const watchdog = page.waitForFunction(
+        async () => window.__FOO === 'hit',
+        {
+          polling: 'raf',
+        }
+      );
+      await page.evaluate(async () => (window.__FOO = 'hit'));
+      await watchdog;
+    });
     itFailsFirefox('should work with strict CSP policy', async () => {
       const { page, server } = getTestState();