From e0f5b9392373b3a1d4214d27a36fb0c0a369792f Mon Sep 17 00:00:00 2001
From: Xingan Wang <wangxgwxg@gmail.com>
Date: Fri, 3 Nov 2017 00:05:38 -0700
Subject: [PATCH] feat(Frame): addStyleTag and addScriptTag now return
 ElementHandles.

This patch teaches `page.addStyleTag` and `page.addScriptTag` to return ElementHandles
to the added tags.

Fixes #1179.
---
 docs/api.md         |  4 ++--
 lib/FrameManager.js | 34 ++++++++++++++++++++--------------
 test/test.js        | 19 ++++++++++++-------
 3 files changed, 34 insertions(+), 23 deletions(-)

diff --git a/docs/api.md b/docs/api.md
index 5d1d0c44573..bc36b2e76a8 100644
--- a/docs/api.md
+++ b/docs/api.md
@@ -1442,7 +1442,7 @@ const html = await frame.$eval('.main-container', e => e.outerHTML);
   - `url` <[string]> Url of a script to be added.
   - `path` <[string]> Path to the JavaScript file to be injected into frame. If `path` is a relative path, then it is resolved relative to [current working directory](https://nodejs.org/api/process.html#process_process_cwd).
   - `content` <[string]> Raw JavaScript content to be injected into frame.
-- returns: <[Promise]> which resolves when the script's onload fires or when the script content was injected into frame.
+- returns: <[Promise]<[ElementHandle]>> which resolves to the added tag when the script's onload fires or when the script content was injected into frame.
 
 Adds a `<script>` tag into the page with the desired url or content.
 
@@ -1451,7 +1451,7 @@ Adds a `<script>` tag into the page with the desired url or content.
   - `url` <[string]> Url of the `<link>` tag.
   - `path` <[string]> Path to the CSS file to be injected into frame. If `path` is a relative path, then it is resolved relative to [current working directory](https://nodejs.org/api/process.html#process_process_cwd).
   - `content` <[string]> Raw CSS content to be injected into frame.
-- returns: <[Promise]> which resolves when the stylesheet's onload fires or when the CSS content was injected into frame.
+- returns: <[Promise]<[ElementHandle]>> which resolves to the added tag when the stylesheet's onload fires or when the CSS content was injected into frame.
 
 Adds a `<link rel="stylesheet">` tag into the page with the desired url or a `<style type="text/css">` tag with the content.
 
diff --git a/lib/FrameManager.js b/lib/FrameManager.js
index 222f1203fe8..578ac9a73fb 100644
--- a/lib/FrameManager.js
+++ b/lib/FrameManager.js
@@ -323,86 +323,92 @@ class Frame {
 
   /**
    * @param {Object} options
-   * @return {!Promise}
+   * @return {!Promise<!ElementHandle>}
    */
   async addScriptTag(options) {
     if (typeof options.url === 'string')
-      return this.evaluate(addScriptUrl, options.url);
+      return this._context.evaluateHandle(addScriptUrl, options.url);
 
     if (typeof options.path === 'string') {
       let contents = await readFileAsync(options.path, 'utf8');
       contents += '//# sourceURL=' + options.path.replace(/\n/g, '');
 
-      return this.evaluate(addScriptContent, contents);
+      return this._context.evaluateHandle(addScriptContent, contents);
     }
 
     if (typeof options.content === 'string')
-      return this.evaluate(addScriptContent, options.content);
+      return this._context.evaluateHandle(addScriptContent, options.content);
 
     throw new Error('Provide an object with a `url`, `path` or `content` property');
 
     /**
      * @param {string} url
+     * @return {!Promise<!HTMLElement>}
      */
-    function addScriptUrl(url) {
+    async function addScriptUrl(url) {
       const script = document.createElement('script');
       script.src = url;
-      const promise = new Promise(x => script.onload = x);
       document.head.appendChild(script);
-      return promise;
+      await new Promise(x => script.onload = x);
+      return script;
     }
 
     /**
      * @param {string} content
+     * @return {!HTMLElement}
      */
     function addScriptContent(content) {
       const script = document.createElement('script');
       script.type = 'text/javascript';
       script.text = content;
       document.head.appendChild(script);
+      return script;
     }
   }
 
   /**
    * @param {Object} options
-   * @return {!Promise}
+   * @return {!Promise<!ElementHandle>}
    */
   async addStyleTag(options) {
     if (typeof options.url === 'string')
-      return this.evaluate(addStyleUrl, options.url);
+      return await this._context.evaluateHandle(addStyleUrl, options.url);
 
     if (typeof options.path === 'string') {
       let contents = await readFileAsync(options.path, 'utf8');
       contents += '/*# sourceURL=' + options.path.replace(/\n/g, '') + '*/';
 
-      return this.evaluate(addStyleContent, contents);
+      return await this._context.evaluateHandle(addStyleContent, contents);
     }
 
     if (typeof options.content === 'string')
-      return this.evaluate(addStyleContent, options.content);
+      return await this._context.evaluateHandle(addStyleContent, options.content);
 
     throw new Error('Provide an object with a `url`, `path` or `content` property');
 
     /**
      * @param {string} url
+     * @return {!Promise<!HTMLElement>}
      */
-    function addStyleUrl(url) {
+    async function addStyleUrl(url) {
       const link = document.createElement('link');
       link.rel = 'stylesheet';
       link.href = url;
-      const promise = new Promise(x => link.onload = x);
       document.head.appendChild(link);
-      return promise;
+      await new Promise(x => link.onload = x);
+      return link;
     }
 
     /**
      * @param {string} content
+     * @return {!HTMLElement}
      */
     function addStyleContent(content) {
       const style = document.createElement('style');
       style.type = 'text/css';
       style.appendChild(document.createTextNode(content));
       document.head.appendChild(style);
+      return style;
     }
   }
 
diff --git a/test/test.js b/test/test.js
index 14468c910c5..c22655d3940 100644
--- a/test/test.js
+++ b/test/test.js
@@ -2375,13 +2375,15 @@ describe('Page', function() {
 
     it('should work with a url', SX(async function() {
       await page.goto(EMPTY_PAGE);
-      await page.addScriptTag({ url: '/injectedfile.js' });
+      const scriptHandle = await page.addScriptTag({ url: '/injectedfile.js' });
+      expect(scriptHandle.asElement()).not.toBeNull();
       expect(await page.evaluate(() => __injected)).toBe(42);
     }));
 
     it('should work with a path', SX(async function() {
       await page.goto(EMPTY_PAGE);
-      await page.addScriptTag({ path: path.join(__dirname, 'assets/injectedfile.js') });
+      const scriptHandle = await page.addScriptTag({ path: path.join(__dirname, 'assets/injectedfile.js') });
+      expect(scriptHandle.asElement()).not.toBeNull();
       expect(await page.evaluate(() => __injected)).toBe(42);
     }));
 
@@ -2394,7 +2396,8 @@ describe('Page', function() {
 
     it('should work with content', SX(async function() {
       await page.goto(EMPTY_PAGE);
-      await page.addScriptTag({ content: 'window.__injected = 35;' });
+      const scriptHandle = await page.addScriptTag({ content: 'window.__injected = 35;' });
+      expect(scriptHandle.asElement()).not.toBeNull();
       expect(await page.evaluate(() => __injected)).toBe(35);
     }));
   });
@@ -2412,13 +2415,15 @@ describe('Page', function() {
 
     it('should work with a url', SX(async function() {
       await page.goto(EMPTY_PAGE);
-      await page.addStyleTag({ url: '/injectedstyle.css' });
+      const styleHandle = await page.addStyleTag({ url: '/injectedstyle.css' });
+      expect(styleHandle.asElement()).not.toBeNull();
       expect(await page.evaluate(`window.getComputedStyle(document.querySelector('body')).getPropertyValue('background-color')`)).toBe('rgb(255, 0, 0)');
     }));
 
     it('should work with a path', SX(async function() {
       await page.goto(EMPTY_PAGE);
-      await page.addStyleTag({ path: path.join(__dirname, 'assets/injectedstyle.css') });
+      const styleHandle = await page.addStyleTag({ path: path.join(__dirname, 'assets/injectedstyle.css') });
+      expect(styleHandle.asElement()).not.toBeNull();
       expect(await page.evaluate(`window.getComputedStyle(document.querySelector('body')).getPropertyValue('background-color')`)).toBe('rgb(255, 0, 0)');
     }));
 
@@ -2428,12 +2433,12 @@ describe('Page', function() {
       const styleHandle = await page.$('style');
       const styleContent = await page.evaluate(style => style.innerHTML, styleHandle);
       expect(styleContent).toContain(path.join('assets', 'injectedstyle.css'));
-      styleHandle.dispose();
     }));
 
     it('should work with content', SX(async function() {
       await page.goto(EMPTY_PAGE);
-      await page.addStyleTag({ content: 'body { background-color: green; }' });
+      const styleHandle = await page.addStyleTag({ content: 'body { background-color: green; }' });
+      expect(styleHandle.asElement()).not.toBeNull();
       expect(await page.evaluate(`window.getComputedStyle(document.querySelector('body')).getPropertyValue('background-color')`)).toBe('rgb(0, 128, 0)');
     }));
   });