From 05b1aca21e5c2f221373d601d3b8ae6b2a984292 Mon Sep 17 00:00:00 2001
From: JoelEinbinder <joel.einbinder@gmail.com>
Date: Thu, 4 Jan 2018 14:49:13 -0800
Subject: [PATCH] feat: support JSHandles for page.waitFor* calls (#1712)

This patch:
- teaches page.waitFor* methods to accept JSHandles
- starts returning JSHandles from page.waitFor* calls.

BREAKING CHANGE: this patch starts allocating `JSHandle`/`ElementHandle` instances for every call to `page.waitFor*` functions. These handles should be disposed manually to avoid memory consumption.

Fixes #1703, fixes #1654, fixes #1724.
---
 docs/api.md         | 34 +++++++++---------
 lib/FrameManager.js | 87 +++++++++++++++++++++++++++++----------------
 test/test.js        | 17 +++++++++
 3 files changed, 90 insertions(+), 48 deletions(-)

diff --git a/docs/api.md b/docs/api.md
index 989b2c78391..f1fe7d7afa2 100644
--- a/docs/api.md
+++ b/docs/api.md
@@ -502,7 +502,7 @@ Shortcut for [page.mainFrame().$$(selector)](#frameselector-1).
 #### page.$$eval(selector, pageFunction[, ...args])
 - `selector` <[string]> A [selector] to query frame for
 - `pageFunction` <[function]> Function to be evaluated in browser context
-- `...args` <...[Serializable]|[ElementHandle]> Arguments to pass to `pageFunction`
+- `...args` <...[Serializable]|[JSHandle]> Arguments to pass to `pageFunction`
 - returns: <[Promise]<[Serializable]>> Promise which resolves to the return value of `pageFunction`
 
 This method runs `document.querySelectorAll` within the page and passes it as the first argument to `pageFunction`.
@@ -517,7 +517,7 @@ const divsCounts = await page.$$eval('div', divs => divs.length);
 #### page.$eval(selector, pageFunction[, ...args])
 - `selector` <[string]> A [selector] to query page for
 - `pageFunction` <[function]> Function to be evaluated in browser context
-- `...args` <...[Serializable]|[ElementHandle]> Arguments to pass to `pageFunction`
+- `...args` <...[Serializable]|[JSHandle]> Arguments to pass to `pageFunction`
 - returns: <[Promise]<[Serializable]>> Promise which resolves to the return value of `pageFunction`
 
 This method runs `document.querySelector` within the page and passes it as the first argument to `pageFunction`. If there's no element matching `selector`, the method throws an error.
@@ -667,7 +667,7 @@ List of all available devices is available in the source code: [DeviceDescriptor
 
 #### page.evaluate(pageFunction, ...args)
 - `pageFunction` <[function]|[string]> Function to be evaluated in the page context
-- `...args` <...[Serializable]|[ElementHandle]> Arguments to pass to `pageFunction`
+- `...args` <...[Serializable]|[JSHandle]> Arguments to pass to `pageFunction`
 - returns: <[Promise]<[Serializable]>> Resolves to the return value of `pageFunction`
 
 If the function, passed to the `page.evaluate`, returns a [Promise], then `page.evaluate` would wait for the promise to resolve and return its value.
@@ -1162,8 +1162,8 @@ This is a shortcut for [page.mainFrame().url()](#frameurl)
 #### page.waitFor(selectorOrFunctionOrTimeout[, options[, ...args]])
 - `selectorOrFunctionOrTimeout` <[string]|[number]|[function]> A [selector], predicate or timeout to wait for
 - `options` <[Object]> Optional waiting parameters
-- `...args` <...[Serializable]> Arguments to pass to  `pageFunction`
-- returns: <[Promise]>
+- `...args` <...[Serializable]|[JSHandle]> Arguments to pass to  `pageFunction`
+- returns: <[Promise]<[JSHandle]>> Promise which resolves to a JSHandle of the success value
 
 This method behaves differently with respect to the type of the first parameter:
 - if `selectorOrFunctionOrTimeout` is a `string`, then the first argument is treated as a [selector] to wait for and the method is a shortcut for [page.waitForSelector](#pagewaitforselectorselector-options)
@@ -1180,8 +1180,8 @@ Shortcut for [page.mainFrame().waitFor(selectorOrFunctionOrTimeout[, options[, .
     - `raf` - to constantly execute `pageFunction` in `requestAnimationFrame` callback. This is the tightest polling mode which is suitable to observe styling changes.
     - `mutation` - to execute `pageFunction` on every DOM mutation.
   - `timeout` <[number]> maximum time to wait for in milliseconds. Defaults to `30000` (30 seconds).
-- `...args` <...[Serializable]> Arguments to pass to  `pageFunction`
-- returns: <[Promise]> Promise which resolves when the `pageFunction` returns a truthy value.
+- `...args` <...[Serializable]|[JSHandle]> Arguments to pass to  `pageFunction`
+- returns: <[Promise]<[JSHandle]>> Promise which resolves when the `pageFunction` returns a truthy value. It resolves to a JSHandle of the truthy value.
 
 The `waitForFunction` can be used to observe viewport size change:
 ```js
@@ -1213,7 +1213,7 @@ Shortcut for [page.mainFrame().waitForFunction(pageFunction[, options[, ...args]
   - `visible` <[boolean]> wait for element to be present in DOM and to be visible, i.e. to not have `display: none` or `visibility: hidden` CSS properties. Defaults to `false`.
   - `hidden` <[boolean]> wait for element to not be found in the DOM or to be hidden, i.e. have `display: none` or `visibility: hidden` CSS properties. Defaults to `false`.
   - `timeout` <[number]> maximum time to wait for in milliseconds. Defaults to `30000` (30 seconds).
-- returns: <[Promise]> Promise which resolves when element specified by selector string is added to DOM.
+- returns: <[Promise]<[ElementHandle]>> Promise which resolves when element specified by selector string is added to DOM.
 
 Wait for the `selector` to appear in page. If at the moment of calling
 the method the `selector` already exists, the method will return
@@ -1492,7 +1492,7 @@ The method runs `document.querySelectorAll` within the frame. If no elements mat
 #### frame.$$eval(selector, pageFunction[, ...args])
 - `selector` <[string]> A [selector] to query frame for
 - `pageFunction` <[function]> Function to be evaluated in browser context
-- `...args` <...[Serializable]|[ElementHandle]> Arguments to pass to `pageFunction`
+- `...args` <...[Serializable]|[JSHandle]> Arguments to pass to `pageFunction`
 - returns: <[Promise]<[Serializable]>> Promise which resolves to the return value of `pageFunction`
 
 This method runs `document.querySelectorAll` within the frame and passes it as the first argument to `pageFunction`.
@@ -1507,7 +1507,7 @@ const divsCounts = await frame.$$eval('div', divs => divs.length);
 #### frame.$eval(selector, pageFunction[, ...args])
 - `selector` <[string]> A [selector] to query frame for
 - `pageFunction` <[function]> Function to be evaluated in browser context
-- `...args` <...[Serializable]|[ElementHandle]> Arguments to pass to `pageFunction`
+- `...args` <...[Serializable]|[JSHandle]> Arguments to pass to `pageFunction`
 - returns: <[Promise]<[Serializable]>> Promise which resolves to the return value of `pageFunction`
 
 This method runs `document.querySelector` within the frame and passes it as the first argument to `pageFunction`. If there's no element matching `selector`, the method throws an error.
@@ -1555,7 +1555,7 @@ Gets the full HTML contents of the frame, including the doctype.
 
 #### frame.evaluate(pageFunction, ...args)
 - `pageFunction` <[function]|[string]> Function to be evaluated in browser context
-- `...args` <...[Serializable]|[ElementHandle]> Arguments to pass to  `pageFunction`
+- `...args` <...[Serializable]|[JSHandle]> Arguments to pass to  `pageFunction`
 - returns: <[Promise]<[Serializable]>> Promise which resolves to function return value
 
 If the function, passed to the `frame.evaluate`, returns a [Promise], then `frame.evaluate` would wait for the promise to resolve and return its value.
@@ -1630,8 +1630,8 @@ Returns frame's url.
 #### frame.waitFor(selectorOrFunctionOrTimeout[, options[, ...args]])
 - `selectorOrFunctionOrTimeout` <[string]|[number]|[function]> A [selector], predicate or timeout to wait for
 - `options` <[Object]> Optional waiting parameters
-- `...args` <...[Serializable]> Arguments to pass to  `pageFunction`
-- returns: <[Promise]>
+- `...args` <...[Serializable]|[JSHandle]> Arguments to pass to  `pageFunction`
+- returns: <[Promise]<[JSHandle]>> Promise which resolves to a JSHandle of the success value
 
 This method behaves differently with respect to the type of the first parameter:
 - if `selectorOrFunctionOrTimeout` is a `string`, then the first argument is treated as a [selector] to wait for and the method is a shortcut for [frame.waitForSelector](#framewaitforselectorselector-options)
@@ -1647,8 +1647,8 @@ This method behaves differently with respect to the type of the first parameter:
     - `raf` - to constantly execute `pageFunction` in `requestAnimationFrame` callback. This is the tightest polling mode which is suitable to observe styling changes.
     - `mutation` - to execute `pageFunction` on every DOM mutation.
   - `timeout` <[number]> maximum time to wait for in milliseconds. Defaults to `30000` (30 seconds).
-- `...args` <...[Serializable]> Arguments to pass to  `pageFunction`
-- returns: <[Promise]> Promise which resolves when the `pageFunction` returns a truthy value.
+- `...args` <...[Serializable]|[JSHandle]> Arguments to pass to  `pageFunction`
+- returns: <[Promise]<[JSHandle]>> Promise which resolves when the `pageFunction` returns a truthy value. It resolves to a JSHandle of the truthy value.
 
 The `waitForFunction` can be used to observe viewport size change:
 ```js
@@ -1669,7 +1669,7 @@ puppeteer.launch().then(async browser => {
   - `visible` <[boolean]> wait for element to be present in DOM and to be visible, i.e. to not have `display: none` or `visibility: hidden` CSS properties. Defaults to `false`.
   - `hidden` <[boolean]> wait for element to not be found in the DOM or to be hidden, i.e. have `display: none` or `visibility: hidden` CSS properties. Defaults to `false`.
   - `timeout` <[number]> maximum time to wait for in milliseconds. Defaults to `30000` (30 seconds).
-- returns: <[Promise]> Promise which resolves when element specified by selector string is added to DOM.
+- returns: <[Promise]<[ElementHandle]>> Promise which resolves when element specified by selector string is added to DOM.
 
 Wait for the `selector` to appear in page. If at the moment of calling
 the method the `selector` already exists, the method will return
@@ -1699,7 +1699,7 @@ The class represents a context for JavaScript execution. Examples of JavaScript
 
 #### executionContext.evaluate(pageFunction, ...args)
 - `pageFunction` <[function]|[string]> Function to be evaluated in `executionContext`
-- `...args` <...[Serializable]|[ElementHandle]> Arguments to pass to `pageFunction`
+- `...args` <...[Serializable]|[JSHandle]> Arguments to pass to `pageFunction`
 - returns: <[Promise]<[Serializable]>> Promise which resolves to function return value
 
 If the function, passed to the `executionContext.evaluate`, returns a [Promise], then `executionContext.evaluate` would wait for the promise to resolve and return its value.
diff --git a/lib/FrameManager.js b/lib/FrameManager.js
index 1a3822b840f..70e2307f113 100644
--- a/lib/FrameManager.js
+++ b/lib/FrameManager.js
@@ -583,17 +583,18 @@ class Frame {
      * @param {string} selector
      * @param {boolean} waitForVisible
      * @param {boolean} waitForHidden
-     * @return {boolean}
+     * @return {?Node|boolean}
      */
     function predicate(selector, waitForVisible, waitForHidden) {
       const node = document.querySelector(selector);
       if (!node)
         return waitForHidden;
       if (!waitForVisible && !waitForHidden)
-        return true;
+        return node;
       const style = window.getComputedStyle(node);
       const isVisible = style && style.visibility !== 'hidden' && hasVisibleBoundingBox();
-      return (waitForVisible === isVisible || waitForHidden === !isVisible);
+      const success = (waitForVisible === isVisible || waitForHidden === !isVisible);
+      return success ? node : null;
 
       /**
        * @return {boolean}
@@ -606,15 +607,14 @@ class Frame {
   }
 
   /**
-   * @param {Function} pageFunction
+   * @param {Function|string} pageFunction
    * @param {!Object=} options
    * @return {!Promise}
    */
   waitForFunction(pageFunction, options = {}, ...args) {
     const timeout = options.timeout || 30000;
     const polling = options.polling || 'raf';
-    const predicateCode = 'return ' + helper.evaluationString(pageFunction, ...args);
-    return new WaitTask(this, predicateCode, polling, timeout).promise;
+    return new WaitTask(this, pageFunction, polling, timeout, ...args).promise;
   }
 
   /**
@@ -658,11 +658,12 @@ helper.tracePublicAPI(Frame);
 class WaitTask {
   /**
    * @param {!Frame} frame
-   * @param {string} predicateBody
+   * @param {Function|string} predicateBody
    * @param {string|number} polling
    * @param {number} timeout
+   * @param {!Array<*>} args
    */
-  constructor(frame, predicateBody, polling, timeout) {
+  constructor(frame, predicateBody, polling, timeout, ...args) {
     if (helper.isString(polling))
       console.assert(polling === 'raf' || polling === 'mutation', 'Unknown polling option: ' + polling);
     else if (helper.isNumber(polling))
@@ -671,7 +672,10 @@ class WaitTask {
       throw new Error('Unknown polling options: ' + polling);
 
     this._frame = frame;
-    this._pageScript = helper.evaluationString(waitForPredicatePageFunction, predicateBody, polling, timeout);
+    this._polling = polling;
+    this._timeout = timeout;
+    this._predicateBody = helper.isString(predicateBody) ? 'return ' + predicateBody : 'return (' + predicateBody + ')(...args)';
+    this._args = args;
     this._runCount = 0;
     frame._waitTasks.add(this);
     this.promise = new Promise((resolve, reject) => {
@@ -695,20 +699,26 @@ class WaitTask {
 
   async rerun() {
     const runCount = ++this._runCount;
-    let success = false;
+    /** @type {?JSHandle} */
+    let success = null;
     let error = null;
     try {
-      success = await this._frame.evaluate(this._pageScript);
+      success = await (await this._frame.executionContext()).evaluateHandle(waitForPredicatePageFunction, this._predicateBody, this._polling, this._timeout, ...this._args);
     } catch (e) {
       error = e;
     }
 
-    if (this._terminated || runCount !== this._runCount)
+    if (this._terminated || runCount !== this._runCount) {
+      if (success)
+        await success.dispose();
       return;
+    }
 
     // Ignore timeouts in pageScript - we track timeouts ourselves.
-    if (!success && !error)
+    if (!error && !(await success.jsonValue())) {
+      await success.dispose();
       return;
+    }
 
     // When the page is navigated, the promise is rejected.
     // We will try again in the new execution context.
@@ -723,7 +733,7 @@ class WaitTask {
     if (error)
       this._reject(error);
     else
-      this._resolve();
+      this._resolve(success);
 
     this._cleanup();
   }
@@ -739,34 +749,39 @@ class WaitTask {
  * @param {string} predicateBody
  * @param {string} polling
  * @param {number} timeout
- * @return {!Promise<boolean>}
+ * @return {!Promise<*>}
  */
-async function waitForPredicatePageFunction(predicateBody, polling, timeout) {
-  const predicate = new Function(predicateBody);
+async function waitForPredicatePageFunction(predicateBody, polling, timeout, ...args) {
+  const predicate = new Function('...args', predicateBody);
   let timedOut = false;
   setTimeout(() => timedOut = true, timeout);
   if (polling === 'raf')
-    await pollRaf();
-  else if (polling === 'mutation')
-    await pollMutation();
-  else if (typeof polling === 'number')
-    await pollInterval(polling);
-  return !timedOut;
+    return await pollRaf();
+  if (polling === 'mutation')
+    return await pollMutation();
+  if (typeof polling === 'number')
+    return await pollInterval(polling);
 
   /**
-   * @return {!Promise}
+   * @return {!Promise<*>}
    */
   function pollMutation() {
-    if (predicate())
-      return Promise.resolve();
+    const success = predicate.apply(null, args);
+    if (success)
+      return Promise.resolve(success);
 
     let fulfill;
     const result = new Promise(x => fulfill = x);
     const observer = new MutationObserver(mutations => {
-      if (timedOut || predicate()) {
+      if (timedOut) {
         observer.disconnect();
         fulfill();
       }
+      const success = predicate.apply(null, args);
+      if (success) {
+        observer.disconnect();
+        fulfill(success);
+      }
     });
     observer.observe(document, {
       childList: true,
@@ -777,7 +792,7 @@ async function waitForPredicatePageFunction(predicateBody, polling, timeout) {
   }
 
   /**
-   * @return {!Promise}
+   * @return {!Promise<*>}
    */
   function pollRaf() {
     let fulfill;
@@ -786,8 +801,13 @@ async function waitForPredicatePageFunction(predicateBody, polling, timeout) {
     return result;
 
     function onRaf() {
-      if (timedOut || predicate())
+      if (timedOut) {
         fulfill();
+        return;
+      }
+      const success = predicate.apply(null, args);
+      if (success)
+        fulfill(success);
       else
         requestAnimationFrame(onRaf);
     }
@@ -795,7 +815,7 @@ async function waitForPredicatePageFunction(predicateBody, polling, timeout) {
 
   /**
    * @param {number} pollInterval
-   * @return {!Promise}
+   * @return {!Promise<*>}
    */
   function pollInterval(pollInterval) {
     let fulfill;
@@ -804,8 +824,13 @@ async function waitForPredicatePageFunction(predicateBody, polling, timeout) {
     return result;
 
     function onTimeout() {
-      if (timedOut || predicate())
+      if (timedOut) {
         fulfill();
+        return;
+      }
+      const success = predicate.apply(null, args);
+      if (success)
+        fulfill(success);
       else
         setTimeout(onTimeout, pollInterval);
     }
diff --git a/test/test.js b/test/test.js
index 746be4d56f4..8d0bace650f 100644
--- a/test/test.js
+++ b/test/test.js
@@ -692,6 +692,18 @@ describe('Page', function() {
       expect(error).toBeTruthy();
       expect(error.message).toContain('Cannot poll with non-positive interval');
     });
+    it('should return the success value as a JSHandle', async({page}) => {
+      expect(await (await page.waitForFunction(() => 5)).jsonValue()).toBe(5);
+    });
+    it('should accept ElementHandle arguments', async({page}) => {
+      await page.setContent('<div></div>');
+      const div = await page.$('div');
+      let resolved = false;
+      const waitForFunction = page.waitForFunction(element => !element.parentElement, {}, div).then(() => resolved = true);
+      expect(resolved).toBe(false);
+      await page.evaluate(element => element.remove(), div);
+      await waitForFunction;
+    });
   });
 
   describe('Frame.waitForSelector', function() {
@@ -858,6 +870,11 @@ describe('Page', function() {
       await page.evaluate(() => document.querySelector('div').className = 'zombo');
       expect(await waitForSelector).toBe(true);
     });
+    it('should return the element handle', async({page, server}) => {
+      const waitForSelector = page.waitForSelector('.zombo');
+      await page.setContent(`<div class='zombo'>anything</div>`);
+      expect(await page.evaluate(x => x.textContent, await waitForSelector)).toBe('anything');
+    });
   });
 
   describe('Page.waitFor', function() {