From 1f954fa7edb277328be3c81f76c4d4b73bc61652 Mon Sep 17 00:00:00 2001
From: Andrey Lushnikov <lushnikov@chromium.org>
Date: Fri, 21 Jul 2017 11:46:57 -0700
Subject: [PATCH] Implement timeout option for page.waitFor

This patch implements timeout option for page.waitFor. The function
will throw if the selector doesn't appear during timeout milliseconds
of waittime.

References #89, #91.
---
 docs/api.md         |  5 +++--
 lib/FrameManager.js | 20 +++++++++++++++-----
 test/test.js        |  6 ++++++
 3 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/docs/api.md b/docs/api.md
index 54735bd2..388ffa03 100644
--- a/docs/api.md
+++ b/docs/api.md
@@ -804,12 +804,13 @@ Returns frame's url.
 #### frame.waitFor(selector[, options])
 - `selector` <[string]> CSS selector of awaited element,
 - `options` <[Object]> Optional waiting parameters
-  - `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.
+  - `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`.
+  - `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.
 
 Wait for the `selector` to appear in page. If at the moment of calling
 the method the `selector` already exists, the method will return
-immediately.
+immediately. If the selector doesn't appear after the `timeout` milliseconds of waiting, the function will throw.
 
 This method works across navigations:
 ```js
diff --git a/lib/FrameManager.js b/lib/FrameManager.js
index 50a62f8a..aa3afc7f 100644
--- a/lib/FrameManager.js
+++ b/lib/FrameManager.js
@@ -174,16 +174,17 @@ class FrameManager extends EventEmitter {
    * @param {!Frame} frame
    * @param {string} selector
    * @param {boolean} waitForVisible
+   * @param {number} timeout
    * @return {!Promise<undefined>}
    */
-  async _waitForSelector(frame, selector, waitForVisible) {
+  async _waitForSelector(frame, selector, waitForVisible, timeout) {
     let contextId = undefined;
     if (!frame.isMainFrame()) {
       contextId = this._frameIdToExecutionContextId.get(frame._id);
       console.assert(contextId, 'Frame does not have default context to evaluate in!');
     }
     let { exceptionDetails } = await this._client.send('Runtime.evaluate', {
-      expression: helper.evaluationString(inPageWatchdog, selector, waitForVisible),
+      expression: helper.evaluationString(inPageWatchdog, selector, waitForVisible, timeout),
       contextId,
       awaitPromise: true,
       returnByValue: false,
@@ -194,10 +195,15 @@ class FrameManager extends EventEmitter {
     /**
      * @param {string} selector
      * @param {boolean} waitForVisible
+     * @param {number} timeout
      * @return {!Promise}
      */
-    function inPageWatchdog(selector, visible) {
-      return visible ? waitForVisible(selector) : waitInDOM(selector);
+    async function inPageWatchdog(selector, visible, timeout) {
+      const resultPromise = visible ? waitForVisible(selector) : waitInDOM(selector);
+      const timeoutPromise = new Promise((resolve, reject) => {
+        setTimeout(reject.bind(null, new Error(`waitFor failed: timeout ${timeout}ms exceeded.`)), timeout);
+      });
+      await Promise.race([resultPromise, timeoutPromise]);
 
       /**
        * @param {string} selector
@@ -342,7 +348,11 @@ class Frame {
    * @return {!Promise}
    */
   async waitFor(selector, options = {}) {
-    const awaitedElement = new AwaitedElement(() => this._frameManager._waitForSelector(this, selector, !!options.visible));
+    const timeout = options.timeout || 30000;
+    const awaitedElement = new AwaitedElement(() => this._frameManager._waitForSelector(this, selector, !!options.visible, timeout));
+    // Since navigation will re-install page watchdogs, we should timeout on our
+    // end as well.
+    setTimeout(() => awaitedElement.terminate(new Error(`waitFor failed: timeout ${timeout}ms exceeded`)), timeout);
 
     this._awaitedElements.add(awaitedElement);
     let cleanup = () => this._awaitedElements.delete(awaitedElement);
diff --git a/test/test.js b/test/test.js
index bd751826..2cb62003 100644
--- a/test/test.js
+++ b/test/test.js
@@ -274,6 +274,12 @@ describe('Puppeteer', function() {
       await page.evaluate(() => document.querySelector('div').style.removeProperty('visibility'));
       expect(await waitFor).toBe(true);
     }));
+    it('should respect timeout', SX(async function() {
+      let error = null;
+      await page.waitFor('div', {timeout: 10}).catch(e => error = e);
+      expect(error).toBeTruthy();
+      expect(error.message).toContain('waitFor failed: timeout');
+    }));
   });
 
   describe('Page.Events.Console', function() {