From 8b09358a56bda6ae9d18f5f516530efc85388229 Mon Sep 17 00:00:00 2001
From: Dmitry Gozman <dgozman@gmail.com>
Date: Mon, 13 Jan 2020 16:29:08 -0800
Subject: [PATCH] fix(screenshot): element screenshot should not throw when
 viewport is null (#472)

---
 src/screenshotter.ts        | 37 ++++++++++++++++++++++++++-----------
 test/browsercontext.spec.js | 27 +++++++++++++++++++++++++++
 2 files changed, 53 insertions(+), 11 deletions(-)

diff --git a/src/screenshotter.ts b/src/screenshotter.ts
index 297c076397..a2a0a950ec 100644
--- a/src/screenshotter.ts
+++ b/src/screenshotter.ts
@@ -96,6 +96,7 @@ export class Screenshotter {
     const rewrittenOptions: types.ScreenshotOptions = { ...options };
     return this._queue.postTask(async () => {
       let overridenViewport: types.Viewport | undefined;
+      let viewportSize: types.Size;
 
       let maybeBoundingBox = await this._page._delegate.getBoundingBoxForScreenshot(handle);
       assert(maybeBoundingBox, 'Node is either not visible or not an HTMLElement');
@@ -104,16 +105,26 @@ export class Screenshotter {
       assert(boundingBox.height !== 0, 'Node has 0 height.');
       boundingBox = enclosingIntRect(boundingBox);
 
-      // TODO: viewport may be null here.
-      const viewport = this._page.viewport()!;
-
+      const viewport = this._page.viewport();
       if (!this._page._delegate.canScreenshotOutsideViewport()) {
-        if (boundingBox.width > viewport.width || boundingBox.height > viewport.height) {
-          overridenViewport = {
-            ...viewport,
-            width: Math.max(viewport.width, boundingBox.width),
-            height: Math.max(viewport.height, boundingBox.height),
-          };
+        if (!viewport) {
+          viewportSize = await this._page.evaluate(() => {
+            if (!document.body || !document.documentElement)
+              return;
+            return {
+              width: Math.max(document.body.offsetWidth, document.documentElement.offsetWidth),
+              height: Math.max(document.body.offsetHeight, document.documentElement.offsetHeight),
+            };
+          });
+          if (!viewportSize)
+            throw new Error(kScreenshotDuringNavigationError);
+        } else {
+          viewportSize = viewport;
+        }
+        if (boundingBox.width > viewportSize.width || boundingBox.height > viewportSize.height) {
+          overridenViewport = {... (viewport || viewportSize)};
+          overridenViewport.width = Math.max(viewportSize.width, boundingBox.width);
+          overridenViewport.height = Math.max(viewportSize.height, boundingBox.height);
           await this._page.setViewport(overridenViewport);
         }
 
@@ -128,8 +139,12 @@ export class Screenshotter {
 
       const result = await this._screenshot(format, rewrittenOptions, overridenViewport || viewport);
 
-      if (overridenViewport)
-        await this._page.setViewport(viewport);
+      if (overridenViewport) {
+        if (viewport)
+          await this._page.setViewport(viewport);
+        else
+          await this._page._delegate.resetViewport(viewportSize!);
+      }
 
       return result;
     }).catch(rewriteError);
diff --git a/test/browsercontext.spec.js b/test/browsercontext.spec.js
index b3f5a003f9..f4c1e4700a 100644
--- a/test/browsercontext.spec.js
+++ b/test/browsercontext.spec.js
@@ -136,6 +136,33 @@ module.exports.describe = function({testRunner, expect, playwright, CHROMIUM, WE
       expect(await page.evaluate('window.innerWidth')).toBe(456);
       expect(await page.evaluate('window.innerHeight')).toBe(789);
     });
+    it('should take element screenshot when default viewport is null and restore back', async({server, newPage}) => {
+      const page = await newPage({ viewport: null });
+      await page.setContent(`
+        <div style="height: 14px">oooo</div>
+        <style>
+        div.to-screenshot {
+          border: 1px solid blue;
+          width: 600px;
+          height: 600px;
+          margin-left: 50px;
+        }
+        ::-webkit-scrollbar{
+          display: none;
+        }
+        </style>
+        <div class="to-screenshot"></div>
+        <div class="to-screenshot"></div>
+        <div class="to-screenshot"></div>
+      `);
+      const sizeBefore = await page.evaluate(() => ({ width: document.body.offsetWidth, height: document.body.offsetHeight }));
+      const elementHandle = await page.$('div.to-screenshot');
+      const screenshot = await elementHandle.screenshot();
+      expect(screenshot).toBeInstanceOf(Buffer);
+      const sizeAfter = await page.evaluate(() => ({ width: document.body.offsetWidth, height: document.body.offsetHeight }));
+      expect(sizeBefore.width).toBe(sizeAfter.width);
+      expect(sizeBefore.height).toBe(sizeAfter.height);
+    });
   });
 
   describe('BrowserContext({setUserAgent})', function() {