Addded builtins object

We can store required methods to this object before any JavaScript is executed and use them later, e.g. in case of waitFor we can store querySelector.

lib/FrameManager.js

@@ -40,6 +40,11 @@ class FrameManager extends EventEmitter {
 
     /** @type {!Map<string, string>} */
     this._frameIdToExecutionContextId = new Map();
+    /** @type {!Set<string>} */
+    this._executionContextIds = new Set();
+
+    /** @type {!Map<string, string>} */
+    this._builtins = new Map();
 
     this._client.on('Page.frameAttached', event => this._onFrameAttached(event.frameId, event.parentFrameId));
     this._client.on('Page.frameNavigated', event => this._onFrameNavigated(event.frame));
@@ -61,6 +66,18 @@ class FrameManager extends EventEmitter {
     return Array.from(this._frames.values());
   }
 
+  /**
+   * @param {string} executionContextId
+   * @param {string} builtinsObjectId
+   */
+  setBuiltins(executionContextId, builtinsObjectId) {
+    if (this._executionContextIds.has(executionContextId)) {
+      this._builtins.set(executionContextId, builtinsObjectId);
+    } else {
+      this._mainFrameBuiltins = builtinsObjectId;
+    }
+  }
+
   /**
    * @param {string} frameId
    * @param {?string} parentFrameId
@@ -104,8 +121,10 @@ class FrameManager extends EventEmitter {
   }
 
   _onExecutionContextCreated(context) {
-    if (context.auxData && context.auxData.isDefault && context.auxData.frameId)
+    if (context.auxData && context.auxData.isDefault && context.auxData.frameId) {
       this._frameIdToExecutionContextId.set(context.auxData.frameId, context.id);
+      this._executionContextIds.add(context.id);
+    }
   }
 
   /**
@@ -118,6 +137,7 @@ class FrameManager extends EventEmitter {
     for (let child of frame.childFrames())
       this._removeFramesRecursively(child);
     this._frames.delete(frame._id, frame);
+    this._executionContextIds.delete(this._frameIdToExecutionContextId.get(frame._id));
     this._frameIdToExecutionContextId.delete(frame._id);
     frame._id = newFrameId;
     frame._adoptPayload(newFramePayload);
@@ -148,6 +168,7 @@ class FrameManager extends EventEmitter {
       this._removeFramesRecursively(child);
     frame._detach();
     this._frames.delete(frame._id);
+    this._executionContextIds.delete(this._frameIdToExecutionContextId.get(frame._id));
     this._frameIdToExecutionContextId.delete(frame._id);
     this.emit(FrameManager.Events.FrameDetached, frame);
   }
@@ -201,34 +222,39 @@ class FrameManager extends EventEmitter {
    * @return {!Promise<undefined>}
    */
   async _waitForInFrame(selector, frame) {
-    let code = selector => new Promise((fulfill, reject) => {
-      if (document.querySelector(selector)) {
-        fulfill();
-        return;
-      }
-      new MutationObserver((mutations, observer) => {
-        for (let mutation of mutations) {
-          for (let node of mutation.addedNodes) {
-            if (node.matches(selector)) {
-              observer.disconnect();
-              fulfill();
-              return;
-            }
+    function code(selector) {
+      return function() {
+        return new Promise((fulfill, reject) => {
+          if (this.querySelector(selector)) {
+            fulfill();
+            return;
           }
-        }
-      }).observe(document.documentElement, {
-        childList: true,
-        subtree: true
-      });
-    });
+          new MutationObserver((mutations, observer) => {
+            for (let mutation of mutations) {
+              for (let node of mutation.addedNodes) {
+                if (node.matches(selector)) {
+                  observer.disconnect();
+                  fulfill();
+                  return;
+                }
+              }
+            }
+          }).observe(document.documentElement, {
+            childList: true,
+            subtree: true
+          });
+        });
+      };
+    }
     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(code, selector),
-      contextId,
+    let objectId = this._builtins.has(contextId) ? this._builtins.get(contextId) : this._mainFrameBuiltins;
+    let { exceptionDetails } = await this._client.send('Runtime.callFunctionOn', {
+      objectId,
+      functionDeclaration: helper.evaluationString(code, selector),
       awaitPromise: true,
       returnByValue: false,
     });

lib/Page.js

@@ -44,6 +44,16 @@ class Page extends EventEmitter {
     let page = new Page(client, frameManager, networkManager, screenDPI);
     // Initialize default page size.
     await page.setViewportSize({width: 400, height: 300});
+    // Initialize builtins.
+    function createBuiltins() {
+      // eslint-disable-next-line no-console
+      console.debug('driver:Builtins', {
+        querySelector: document.querySelector.bind(document),
+        __proto__: null
+      });
+    }
+    await page.evaluate(createBuiltins);
+    await page.evaluateOnInitialized(createBuiltins);
     return page;
   }
 
@@ -216,6 +226,10 @@ class Page extends EventEmitter {
       }
       return;
     }
+    if (event.type === 'debug' && event.args.length > 1 && event.args[0].value === 'driver:Builtins') {
+      this._frameManager.setBuiltins(event.executionContextId, event.args[1].objectId);
+      return;
+    }
     let values = event.args.map(arg => arg.value || arg.description || '');
     this.emit(Page.Events.ConsoleMessage, values.join(' '));
   }

test/test.js

@@ -212,16 +212,28 @@ describe('Puppeteer', function() {
       expect(added).toBe(true);
     }));
 
-    it('should throw if evaluation failed', SX(async function() {
+    it('should not throw if document.querySelector is redefined', SX(async function() {
       await page.evaluateOnInitialized(function() {
         document.querySelector = null;
       });
       await page.navigate(EMPTY_PAGE);
       try {
         await page.waitFor('*');
+      } catch (e) {
+        fail('Failed waitFor threw.');
+      }
+    }));
+
+    it('should throw if evaluation failed', SX(async function() {
+      await page.evaluateOnInitialized(function() {
+        window.MutationObserver = null;
+      });
+      await page.navigate(EMPTY_PAGE);
+      try {
+        await page.waitFor('div');
         fail('Failed waitFor did not throw.');
       } catch (e) {
-        expect(e.message).toBe('Evaluation failed: document.querySelector is not a function');
+        expect(e.message).toBe('Evaluation failed: MutationObserver is not a constructor');
       }
     }));
   });