From ae22ef30b38673d1d09fb7cf2f7b7bffd3517f5a Mon Sep 17 00:00:00 2001
From: Andrey Lushnikov <aslushnikov@gmail.com>
Date: Tue, 13 Feb 2018 14:02:44 -0800
Subject: [PATCH] feat(ExecutionContext): introduce ExecutionContext.frame()
 (#1972)

This patch introduces ExecutionContext.frame() that returns Frame
associated with this Execution Context.

This allows to associate console messages with the originating frame,
if any.
---
 docs/api.md             |  6 ++++++
 lib/ExecutionContext.js | 15 ++++++++++-----
 lib/FrameManager.js     | 13 ++++++-------
 test/test.js            |  5 ++++-
 4 files changed, 26 insertions(+), 13 deletions(-)

diff --git a/docs/api.md b/docs/api.md
index 3731c4e6..fc30975a 100644
--- a/docs/api.md
+++ b/docs/api.md
@@ -167,6 +167,7 @@
 - [class: ExecutionContext](#class-executioncontext)
   * [executionContext.evaluate(pageFunction, ...args)](#executioncontextevaluatepagefunction-args)
   * [executionContext.evaluateHandle(pageFunction, ...args)](#executioncontextevaluatehandlepagefunction-args)
+  * [executionContext.frame()](#executioncontextframe)
   * [executionContext.queryObjects(prototypeHandle)](#executioncontextqueryobjectsprototypehandle)
 - [class: JSHandle](#class-jshandle)
   * [jsHandle.asElement()](#jshandleaselement)
@@ -2041,6 +2042,11 @@ await aHandle.dispose();
 await resultHandle.dispose();
 ```
 
+#### executionContext.frame()
+- returns: <?[Frame]> Frame associated with this execution context.
+
+> **NOTE** Not every execution context is associated with a frame. For example, workers and extensions have execution contexts that are not associated with frames.
+
 #### executionContext.queryObjects(prototypeHandle)
 - `prototypeHandle` <[JSHandle]> A handle to the object prototype.
 - returns: <[JSHandle]> A handle to an array of objects with this prototype
diff --git a/lib/ExecutionContext.js b/lib/ExecutionContext.js
index 695bf82c..ff313be3 100644
--- a/lib/ExecutionContext.js
+++ b/lib/ExecutionContext.js
@@ -21,17 +21,22 @@ class ExecutionContext {
    * @param {!Puppeteer.CDPSession} client
    * @param {!Object} contextPayload
    * @param {function(*):!JSHandle} objectHandleFactory
+   * @param {?Puppeteer.Frame} frame
    */
-  constructor(client, contextPayload, objectHandleFactory) {
+  constructor(client, contextPayload, objectHandleFactory, frame) {
     this._client = client;
+    this._frame = frame;
     this._contextId = contextPayload.id;
-
-    const auxData = contextPayload.auxData || {isDefault: true};
-    this._frameId = auxData.frameId || null;
-    this._isDefault = !!auxData.isDefault;
     this._objectHandleFactory = objectHandleFactory;
   }
 
+  /**
+   * @return {?Puppeteer.Frame}
+   */
+  frame() {
+    return this._frame;
+  }
+
   /**
    * @param {Function|string} pageFunction
    * @param {...*} args
diff --git a/lib/FrameManager.js b/lib/FrameManager.js
index d700e0b0..4e542803 100644
--- a/lib/FrameManager.js
+++ b/lib/FrameManager.js
@@ -154,11 +154,11 @@ class FrameManager extends EventEmitter {
   }
 
   _onExecutionContextCreated(contextPayload) {
-    const context = new ExecutionContext(this._client, contextPayload, this.createJSHandle.bind(this, contextPayload.id));
+    const frameId = contextPayload.auxData && contextPayload.auxData.isDefault ? contextPayload.auxData.frameId : null;
+    const frame = frameId ? this._frames.get(frameId) : null;
+    const context = new ExecutionContext(this._client, contextPayload, this.createJSHandle.bind(this, contextPayload.id), frame);
     this._contextIdToContext.set(contextPayload.id, context);
-
-    const frame = context._frameId ? this._frames.get(context._frameId) : null;
-    if (frame && context._isDefault)
+    if (frame)
       frame._setDefaultContext(context);
   }
 
@@ -166,9 +166,8 @@ class FrameManager extends EventEmitter {
    * @param {!ExecutionContext} context
    */
   _removeContext(context) {
-    const frame = context._frameId ? this._frames.get(context._frameId) : null;
-    if (frame && context._isDefault)
-      frame._setDefaultContext(null);
+    if (context.frame())
+      context.frame()._setDefaultContext(null);
   }
 
   /**
diff --git a/test/test.js b/test/test.js
index d2a43175..c3e98d03 100644
--- a/test/test.js
+++ b/test/test.js
@@ -152,7 +152,8 @@ describe('Puppeteer', function() {
       expect(response.ok()).toBe(true);
       expect(response.securityDetails()).toBeTruthy();
       expect(response.securityDetails().protocol()).toBe('TLS 1.2');
-      browser.close();
+      await page.close();
+      await browser.close();
     });
     it('should reject all promises when browser is closed', async() => {
       const browser = await puppeteer.launch(defaultBrowserOptions);
@@ -647,6 +648,8 @@ describe('Page', function() {
       expect(context1).toBeTruthy();
       expect(context2).toBeTruthy();
       expect(context1 !== context2).toBeTruthy();
+      expect(context1.frame()).toBe(frame1);
+      expect(context2.frame()).toBe(frame2);
 
       await Promise.all([
         context1.evaluate(() => window.a = 1),