From 6318ba6e4afd3da0e529e1843f14de3e9c300b5f Mon Sep 17 00:00:00 2001
From: Dmitry Gozman <dgozman@gmail.com>
Date: Wed, 5 Feb 2020 17:20:23 -0800
Subject: [PATCH] feat(frame): introduce frame.frameElement (#856)

Fixes #839.
---
 docs/api.md            | 14 ++++++++++++++
 src/chromium/crPage.ts | 12 ++++++++++++
 src/firefox/ffPage.ts  | 17 +++++++++++++++++
 src/frames.ts          |  4 ++++
 src/page.ts            |  1 +
 src/webkit/wkPage.ts   | 17 +++++++++++++++++
 test/frame.spec.js     | 30 ++++++++++++++++++++++++++++++
 7 files changed, 95 insertions(+)

diff --git a/docs/api.md b/docs/api.md
index 935670dedd1e7..e399429e2f84a 100644
--- a/docs/api.md
+++ b/docs/api.md
@@ -1676,6 +1676,7 @@ An example of getting text from an iframe element:
 - [frame.evaluateHandle(pageFunction[, ...args])](#frameevaluatehandlepagefunction-args)
 - [frame.fill(selector, value, options)](#framefillselector-value-options)
 - [frame.focus(selector, options)](#framefocusselector-options)
+- [frame.frameElement()](#frameframeelement)
 - [frame.goto(url[, options])](#framegotourl-options)
 - [frame.hover(selector[, options])](#framehoverselector-options)
 - [frame.isDetached()](#frameisdetached)
@@ -1916,6 +1917,19 @@ If there's no text `<input>`, `<textarea>` or `[contenteditable]` element matchi
 This method fetches an element with `selector` and focuses it.
 If there's no element matching `selector`, the method throws an error.
 
+#### frame.frameElement()
+- returns: <[Promise]<[ElementHandle]>> Promise that resolves with a `frame` or `iframe` element handle which corresponds to this frame.
+
+This is an inverse of [elementHandle.contentFrame()](#elementhandlecontentframe). Note that returned handle actually belongs to the parent frame.
+
+This method throws an error if the frame has been detached before `frameElement()` returns.
+
+```js
+const frameElement = await frame.frameElement();
+const contentFrame = await frameElement.contentFrame();
+console.log(frame === contentFrame);  // -> true
+```
+
 #### frame.goto(url[, options])
 - `url` <[string]> URL to navigate frame to. The url should include scheme, e.g. `https://`.
 - `options` <[Object]> Navigation parameters which might have the following properties:
diff --git a/src/chromium/crPage.ts b/src/chromium/crPage.ts
index 3b9ebdd156d6d..83e125087cf76 100644
--- a/src/chromium/crPage.ts
+++ b/src/chromium/crPage.ts
@@ -525,6 +525,18 @@ export class CRPage implements PageDelegate {
   coverage(): Coverage | undefined {
     return this._coverage;
   }
+
+  async getFrameElement(frame: frames.Frame): Promise<dom.ElementHandle> {
+    const { backendNodeId } = await this._client.send('DOM.getFrameOwner', { frameId: frame._id }).catch(e => {
+      if (e instanceof Error && e.message.includes('Frame with the given id was not found.'))
+        e.message = 'Frame has been detached.';
+      throw e;
+    });
+    const parent = frame.parentFrame();
+    if (!parent)
+      throw new Error('Frame has been detached.');
+    return this.adoptBackendNodeId(backendNodeId, await parent._mainContext());
+  }
 }
 
 function toRemoteObject(handle: js.JSHandle): Protocol.Runtime.RemoteObject {
diff --git a/src/firefox/ffPage.ts b/src/firefox/ffPage.ts
index 847a85a9f3638..0860e76f79e76 100644
--- a/src/firefox/ffPage.ts
+++ b/src/firefox/ffPage.ts
@@ -443,6 +443,23 @@ export class FFPage implements PageDelegate {
   coverage(): Coverage | undefined {
     return undefined;
   }
+
+  async getFrameElement(frame: frames.Frame): Promise<dom.ElementHandle> {
+    const parent = frame.parentFrame();
+    if (!parent)
+      throw new Error('Frame has been detached.');
+    const context = await parent._utilityContext();
+    const handles = await context._$$('iframe');
+    const items = await Promise.all(handles.map(async handle => {
+      const frame = await handle.contentFrame().catch(e => null);
+      return { handle, frame };
+    }));
+    const result = items.find(item => item.frame === frame);
+    await Promise.all(items.map(item => item === result ? Promise.resolve() : item.handle.dispose()));
+    if (!result)
+      throw new Error('Frame has been detached.');
+    return result.handle;
+  }
 }
 
 export function normalizeWaitUntil(waitUntil: frames.LifecycleEvent | frames.LifecycleEvent[]): frames.LifecycleEvent[] {
diff --git a/src/frames.ts b/src/frames.ts
index 44646ace0fb87..90d76616bcc8d 100644
--- a/src/frames.ts
+++ b/src/frames.ts
@@ -437,6 +437,10 @@ export class Frame {
       throw error;
   }
 
+  async frameElement(): Promise<dom.ElementHandle> {
+    return this._page._delegate.getFrameElement(this);
+  }
+
   _context(contextType: ContextType): Promise<dom.FrameExecutionContext> {
     if (this._detached)
       throw new Error(`Execution Context is not available in detached frame "${this.url()}" (are you trying to evaluate?)`);
diff --git a/src/page.ts b/src/page.ts
index 32de48a3bc283..62c64cd73110b 100644
--- a/src/page.ts
+++ b/src/page.ts
@@ -68,6 +68,7 @@ export interface PageDelegate {
   layoutViewport(): Promise<{ width: number, height: number }>;
   setInputFiles(handle: dom.ElementHandle<HTMLInputElement>, files: types.FilePayload[]): Promise<void>;
   getBoundingBox(handle: dom.ElementHandle): Promise<types.Rect | null>;
+  getFrameElement(frame: frames.Frame): Promise<dom.ElementHandle>;
 
   getAccessibilityTree(needle?: dom.ElementHandle): Promise<{tree: accessibility.AXNode, needle: accessibility.AXNode | null}>;
   pdf?: (options?: types.PDFOptions) => Promise<platform.BufferType>;
diff --git a/src/webkit/wkPage.ts b/src/webkit/wkPage.ts
index cf3b79fd34538..ed669ff4e98dd 100644
--- a/src/webkit/wkPage.ts
+++ b/src/webkit/wkPage.ts
@@ -575,6 +575,23 @@ export class WKPage implements PageDelegate {
     return undefined;
   }
 
+  async getFrameElement(frame: frames.Frame): Promise<dom.ElementHandle> {
+    const parent = frame.parentFrame();
+    if (!parent)
+      throw new Error('Frame has been detached.');
+    const context = await parent._utilityContext();
+    const handles = await context._$$('iframe');
+    const items = await Promise.all(handles.map(async handle => {
+      const frame = await handle.contentFrame().catch(e => null);
+      return { handle, frame };
+    }));
+    const result = items.find(item => item.frame === frame);
+    await Promise.all(items.map(item => item === result ? Promise.resolve() : item.handle.dispose()));
+    if (!result)
+      throw new Error('Frame has been detached.');
+    return result.handle;
+  }
+
   _onRequestWillBeSent(session: WKSession, event: Protocol.Network.requestWillBeSentPayload) {
     if (event.request.url.startsWith('data:'))
       return;
diff --git a/test/frame.spec.js b/test/frame.spec.js
index d99d3f719ad73..6cb912ba81d10 100644
--- a/test/frame.spec.js
+++ b/test/frame.spec.js
@@ -31,6 +31,36 @@ module.exports.describe = function({testRunner, expect, FFOX, CHROMIUM, WEBKIT})
     });
   });
 
+  describe('Frame.frameElement', function() {
+    it('should work', async({page, server}) => {
+      await page.goto(server.EMPTY_PAGE);
+      const frame1 = await utils.attachFrame(page, 'frame1', server.EMPTY_PAGE);
+      const frame2 = await utils.attachFrame(page, 'frame2', server.EMPTY_PAGE);
+      const frame3 = await utils.attachFrame(page, 'frame3', server.EMPTY_PAGE);
+      const frame1handle1 = await page.$('#frame1');
+      const frame1handle2 = await frame1.frameElement();
+      const frame3handle1 = await page.$('#frame3');
+      const frame3handle2 = await frame3.frameElement();
+      expect(await frame1handle1.evaluate((a, b) => a === b, frame1handle2)).toBe(true);
+      expect(await frame3handle1.evaluate((a, b) => a === b, frame3handle2)).toBe(true);
+      expect(await frame1handle1.evaluate((a, b) => a === b, frame3handle1)).toBe(false);
+    });
+    it('should work with contentFrame', async({page, server}) => {
+      await page.goto(server.EMPTY_PAGE);
+      const frame = await utils.attachFrame(page, 'frame1', server.EMPTY_PAGE);
+      const handle = await frame.frameElement();
+      const contentFrame = await handle.contentFrame();
+      expect(contentFrame).toBe(frame);
+    });
+    it('should throw when detached', async({page, server}) => {
+      await page.goto(server.EMPTY_PAGE);
+      const frame1 = await utils.attachFrame(page, 'frame1', server.EMPTY_PAGE);
+      await page.$eval('#frame1', e => e.remove());
+      const error = await frame1.frameElement().catch(e => e);
+      expect(error.message).toBe('Frame has been detached.');
+    });
+  });
+
   describe('Frame.evaluate', function() {
     it('should throw for detached frames', async({page, server}) => {
       const frame1 = await utils.attachFrame(page, 'frame1', server.EMPTY_PAGE);