From 55cfff38c6290bdd2a54b43875c890b142d09031 Mon Sep 17 00:00:00 2001
From: Dmitry Gozman <dgozman@gmail.com>
Date: Mon, 8 Jun 2020 16:29:29 -0700
Subject: [PATCH] fix(waitForFunction): handle predicate that throws (#2488)

Currently, we fail when the predicate throws on the first call,
and timeout when it fails on any other call.

There are two possible ways to handle throwing predicates:
- Fail waitForFunction if predicate throws once. This is good
  since it gives you the error faster.
- Tolerate predicate exceptions. This is good because you do
  not have to worry about non-initialized state during load.

This change implements the former.
---
 src/injected/injectedScript.ts | 35 ++++++++++++++++++++++------------
 test/waittask.spec.js          | 17 +++++++++++++++++
 2 files changed, 40 insertions(+), 12 deletions(-)

diff --git a/src/injected/injectedScript.ts b/src/injected/injectedScript.ts
index 2c480188ac..9b4cbd01b7 100644
--- a/src/injected/injectedScript.ts
+++ b/src/injected/injectedScript.ts
@@ -106,16 +106,21 @@ export default class InjectedScript {
 
   private _pollRaf<T>(progress: types.InjectedScriptProgress, predicate: Predicate<T>): Promise<T> {
     let fulfill: (result: T) => void;
-    const result = new Promise<T>(x => fulfill = x);
+    let reject: (error: Error) => void;
+    const result = new Promise<T>((f, r) => { fulfill = f; reject = r; });
 
     const onRaf = () => {
       if (progress.canceled)
         return;
-      const success = predicate(progress);
-      if (success)
-        fulfill(success);
-      else
-        requestAnimationFrame(onRaf);
+      try {
+        const success = predicate(progress);
+        if (success)
+          fulfill(success);
+        else
+          requestAnimationFrame(onRaf);
+      } catch (e) {
+        reject(e);
+      }
     };
 
     onRaf();
@@ -124,15 +129,21 @@ export default class InjectedScript {
 
   private _pollInterval<T>(progress: types.InjectedScriptProgress, pollInterval: number, predicate: Predicate<T>): Promise<T> {
     let fulfill: (result: T) => void;
-    const result = new Promise<T>(x => fulfill = x);
+    let reject: (error: Error) => void;
+    const result = new Promise<T>((f, r) => { fulfill = f; reject = r; });
+
     const onTimeout = () => {
       if (progress.canceled)
         return;
-      const success = predicate(progress);
-      if (success)
-        fulfill(success);
-      else
-        setTimeout(onTimeout, pollInterval);
+      try {
+        const success = predicate(progress);
+        if (success)
+          fulfill(success);
+        else
+          setTimeout(onTimeout, pollInterval);
+      } catch (e) {
+        reject(e);
+      }
     };
 
     onTimeout();
diff --git a/test/waittask.spec.js b/test/waittask.spec.js
index 0d7624ce7a..5f3a4dddaa 100644
--- a/test/waittask.spec.js
+++ b/test/waittask.spec.js
@@ -66,6 +66,23 @@ describe('Frame.waitForFunction', function() {
     await page.evaluate(() => window.__FOO = 'hit');
     await watchdog;
   });
+  it('should fail with predicate throwing on first call', async({page, server}) => {
+    const error = await page.waitForFunction(() => { throw new Error('oh my'); }).catch(e => e);
+    expect(error.message).toContain('oh my');
+  });
+  it('should fail with predicate throwing sometimes', async({page, server}) => {
+    const error = await page.waitForFunction(() => {
+      window.counter = (window.counter || 0) + 1;
+      if (window.counter === 3)
+        throw new Error('Bad counter!');
+      return window.counter === 5 ? 'result' : false;
+    }).catch(e => e);
+    expect(error.message).toContain('Bad counter!');
+  });
+  it('should fail with ReferenceError on wrong page', async({page, server}) => {
+    const error = await page.waitForFunction(() => globalVar === 123).catch(e => e);
+    expect(error.message).toContain('globalVar');
+  });
   it('should work with strict CSP policy', async({page, server}) => {
     server.setCSP('/empty.html', 'script-src ' + server.PREFIX);
     await page.goto(server.EMPTY_PAGE);