feat(screenshot): accept timeout, migrate to Progress, wait for visible (#2679)

Element screenshot now waits for the element to become visible and
throws on detach.

Both screenshot methods accept a timeout and capture logs using Progress.
Also, carefully handling exceptions and restoring the viewport.

docs/actionability.md

@@ -8,13 +8,11 @@ Some actions like `page.click()` support `{force: true}` option that disable non
 
 | Actions | Performed checks |
 | ------ | ------- |
-| `check()`<br>`click()`<br>`dblclick()`<br>`hover()`<br>`uncheck()` | [Visible]<br>[Stable]<br>[Enabled]<br>[Receiving Events]<br>[Attached]† |
+| `check()`<br>`click()`<br>`dblclick()`<br>`hover()`<br>`uncheck()` | [Visible]<br>[Stable]<br>[Enabled]<br>[Receiving Events]<br>[Attached] |
-| `fill()` | [Visible]<br>[Enabled]<br>[Editable]<br>[Attached]† |
+| `fill()` | [Visible]<br>[Enabled]<br>[Editable]<br>[Attached] |
-| `dispatchEvent()`<br>`focus()`<br>`press()`<br>`setInputFiles()`<br>`selectOption()`<br>`type()` | [Attached]† |
+| `dispatchEvent()`<br>`focus()`<br>`press()`<br>`setInputFiles()`<br>`selectOption()`<br>`type()` | [Attached] |
-| `selectText()`<br>`scrollIntoViewIfNeeded()` | [Visible] |
+| `selectText()`<br>`scrollIntoViewIfNeeded()`<br>`screenshot()` | [Visible]<br>[Attached] |
-| `getAttribute()`<br>`innerText()`<br>`innerHTML()`<br>`textContent()` | [Attached]† |
+| `getAttribute()`<br>`innerText()`<br>`innerHTML()`<br>`textContent()` | [Attached] |
-
-† [Attached] check is only performed during selector-based actions.
 
 ### Visible
 
@@ -40,7 +38,9 @@ Element is considered receiving pointer events when it is the hit target of the
 
 Element is considered attached when it is [connected](https://developer.mozilla.org/en-US/docs/Web/API/Node/isConnected) to a Document or a ShadowRoot.
 
-Attached check is performed during a selector-based action, like `page.click(selector, options)` as opposite to `elementHandle.click(options)`. First, Playwright waits for an element matching `selector` to be attached to the DOM, and then checks that element is still attached before performing the action.
+Attached check differs between selector-based and handle-based actions, like `page.click(selector, options)` as opposite to `elementHandle.click(options)`:
+- For selector-based actions, Playwright first waits for an element matching `selector` to be attached to the DOM, and then checks that element is still attached before performing the action. If element was detached, the action is retried from the start.
+- For handle-based actions, Playwright throws if the element is not attached.
 
 For example, consider a scenario where Playwright will click `Sign Up` button regardless of when the `page.click()` call was made:
 - page is checking that user name is unique and `Sign Up` button is disabled;

docs/api.md

@@ -1610,6 +1610,7 @@ Page routes take precedence over browser context routes (set up with [browserCon
     - `width` <[number]> width of clipping area
     - `height` <[number]> height of clipping area
   - `omitBackground` <[boolean]> Hides default white background and allows capturing screenshots with transparency. Not applicable to `jpeg` images. Defaults to `false`.
+  - `timeout` <[number]> Maximum time in milliseconds, defaults to 30 seconds, pass `0` to disable timeout. The default value can be changed by using the [browserContext.setDefaultTimeout(timeout)](#browsercontextsetdefaulttimeouttimeout) or [page.setDefaultTimeout(timeout)](#pagesetdefaulttimeouttimeout) methods.
 - returns: <[Promise]<[Buffer]>> Promise which resolves to buffer with the captured screenshot.
 
 > **NOTE** Screenshots take at least 1/6 second on Chromium OS X and Chromium Windows. See https://crbug.com/741689 for discussion.
@@ -2856,9 +2857,10 @@ Shortcuts such as `key: "Control+o"` or `key: "Control+Shift+T"` are supported a
   - `type` <"png"|"jpeg"> Specify screenshot type, defaults to `png`.
   - `quality` <[number]> The quality of the image, between 0-100. Not applicable to `png` images.
   - `omitBackground` <[boolean]> Hides default white background and allows capturing screenshots with transparency. Not applicable to `jpeg` images. Defaults to `false`.
+  - `timeout` <[number]> Maximum time in milliseconds, defaults to 30 seconds, pass `0` to disable timeout. The default value can be changed by using the [browserContext.setDefaultTimeout(timeout)](#browsercontextsetdefaulttimeouttimeout) or [page.setDefaultTimeout(timeout)](#pagesetdefaulttimeouttimeout) methods.
 - returns: <[Promise]<[Buffer]>> Promise which resolves to buffer with the captured screenshot.
 
-This method scrolls element into view if needed before taking a screenshot. If the element is detached from DOM, the method throws an error.
+This method waits for the [actionability](./actionability.md) checks, then scrolls element into view before taking a screenshot. If the element is detached from DOM, the method throws an error.
 
 #### elementHandle.scrollIntoViewIfNeeded([options])
 - `options` <[Object]>

src/dom.ts

@@ -202,21 +202,25 @@ export class ElementHandle<T extends Node = Node> extends js.JSHandle<T> {
     return await this._page._delegate.scrollRectIntoViewIfNeeded(this, rect);
   }
 
-  async scrollIntoViewIfNeeded(options: types.TimeoutOptions = {}) {
+  async _waitAndScrollIntoViewIfNeeded(progress: Progress): Promise<void> {
-    return this._runAbortableTask(async progress => {
+    while (progress.isRunning()) {
-      while (progress.isRunning()) {
+      const waited = await this._waitForVisible(progress);
-        const waited = await this._waitForVisible(progress);
+      throwIfNotConnected(waited);
-        throwIfNotConnected(waited);
 
-        progress.throwIfAborted();  // Avoid action that has side-effects.
+      progress.throwIfAborted();  // Avoid action that has side-effects.
-        const result = await this._scrollRectIntoViewIfNeeded();
+      const result = await this._scrollRectIntoViewIfNeeded();
-        throwIfNotConnected(result);
+      throwIfNotConnected(result);
-        if (result === 'notvisible')
+      if (result === 'notvisible')
-          continue;
+        continue;
-        assert(result === 'done');
+      assert(result === 'done');
-        return;
+      return;
-      }
+    }
-    }, this._page._timeoutSettings.timeout(options), 'scrollIntoViewIfNeeded');
+  }
+
+  async scrollIntoViewIfNeeded(options: types.TimeoutOptions = {}) {
+    return this._runAbortableTask(
+        progress => this._waitAndScrollIntoViewIfNeeded(progress),
+        this._page._timeoutSettings.timeout(options), 'scrollIntoViewIfNeeded');
   }
 
   private async _waitForVisible(progress: Progress): Promise<'notconnected' | 'done'> {
@@ -585,8 +589,10 @@ export class ElementHandle<T extends Node = Node> extends js.JSHandle<T> {
     return this._page._delegate.getBoundingBox(this);
   }
 
-  async screenshot(options?: types.ElementScreenshotOptions): Promise<Buffer> {
+  async screenshot(options: types.ElementScreenshotOptions = {}): Promise<Buffer> {
-    return this._page._screenshotter.screenshotElement(this, options);
+    return this._runAbortableTask(
+        progress => this._page._screenshotter.screenshotElement(progress, this, options),
+        this._page._timeoutSettings.timeout(options), 'screenshot');
   }
 
   async $(selector: string): Promise<ElementHandle | null> {

src/page.ts

@@ -436,8 +436,9 @@ export class Page extends EventEmitter {
     return false;
   }
 
-  async screenshot(options?: types.ScreenshotOptions): Promise<Buffer> {
+  async screenshot(options: types.ScreenshotOptions = {}): Promise<Buffer> {
-    return this._screenshotter.screenshotPage(options);
+    const controller = new ProgressController(this._logger, this._timeoutSettings.timeout(options), 'page.screenshot');
+    return controller.run(progress => this._screenshotter.screenshotPage(progress, options));
   }
 
   async title(): Promise<string> {

src/screenshotter.ts

@@ -23,6 +23,7 @@ import { assert, helper } from './helper';
 import { Page } from './page';
 import * as types from './types';
 import { rewriteErrorMessage } from './utils/stackTrace';
+import { Progress } from './progress';
 
 export class Screenshotter {
   private _queue = new TaskQueue();
@@ -66,7 +67,7 @@ export class Screenshotter {
     return fullPageSize;
   }
 
-  async screenshotPage(options: types.ScreenshotOptions = {}): Promise<Buffer> {
+  async screenshotPage(progress: Progress, options: types.ScreenshotOptions): Promise<Buffer> {
     const format = validateScreenshotOptions(options);
     return this._queue.postTask(async () => {
       const { viewportSize, originalViewportSize } = await this._originalViewportSize();
@@ -78,27 +79,30 @@ export class Screenshotter {
         const fitsViewport = fullPageSize.width <= viewportSize.width && fullPageSize.height <= viewportSize.height;
         if (!this._page._delegate.canScreenshotOutsideViewport() && !fitsViewport) {
           overridenViewportSize = fullPageSize;
+          progress.throwIfAborted(); // Avoid side effects.
           await this._page.setViewportSize(overridenViewportSize);
+          progress.cleanupWhenAborted(() => this._restoreViewport(originalViewportSize));
         }
         if (options.clip)
           documentRect = trimClipToSize(options.clip, documentRect);
-        return await this._screenshot(format, documentRect, undefined, options, overridenViewportSize, originalViewportSize);
+        const buffer = await this._screenshot(progress, format, documentRect, undefined, options);
+        progress.throwIfAborted(); // Avoid restoring after failure - should be done by cleanup.
+        if (overridenViewportSize)
+          await this._restoreViewport(originalViewportSize);
+        return buffer;
       }
 
       const viewportRect = options.clip ? trimClipToSize(options.clip, viewportSize) : { x: 0, y: 0, ...viewportSize };
-      return await this._screenshot(format, undefined, viewportRect, options, null, originalViewportSize);
+      return await this._screenshot(progress, format, undefined, viewportRect, options);
     }).catch(rewriteError);
   }
 
-  async screenshotElement(handle: dom.ElementHandle, options: types.ElementScreenshotOptions = {}): Promise<Buffer> {
+  async screenshotElement(progress: Progress, handle: dom.ElementHandle, options: types.ElementScreenshotOptions = {}): Promise<Buffer> {
     const format = validateScreenshotOptions(options);
     return this._queue.postTask(async () => {
       const { viewportSize, originalViewportSize } = await this._originalViewportSize();
 
-      // TODO: make screenshot wait visible, migrate to progress.
+      await handle._waitAndScrollIntoViewIfNeeded(progress);
-      const scrolled = await handle._scrollRectIntoViewIfNeeded();
-      if (scrolled === 'notconnected')
-        throw new Error('Element is not attached to the DOM');
       let boundingBox = await handle.boundingBox();
       assert(boundingBox, 'Node is either not visible or not an HTMLElement');
       assert(boundingBox.width !== 0, 'Node has 0 width.');
@@ -111,11 +115,11 @@ export class Screenshotter {
           width: Math.max(viewportSize.width, boundingBox.width),
           height: Math.max(viewportSize.height, boundingBox.height),
         });
+        progress.throwIfAborted(); // Avoid side effects.
         await this._page.setViewportSize(overridenViewportSize);
+        progress.cleanupWhenAborted(() => this._restoreViewport(originalViewportSize));
 
-        const scrolled = await handle._scrollRectIntoViewIfNeeded();
+        await handle._waitAndScrollIntoViewIfNeeded(progress);
-        if (scrolled === 'notconnected')
-          throw new Error('Element is not attached to the DOM');
         boundingBox = await handle.boundingBox();
         assert(boundingBox, 'Node is either not visible or not an HTMLElement');
         assert(boundingBox.width !== 0, 'Node has 0 width.');
@@ -127,28 +131,42 @@ export class Screenshotter {
       const documentRect = { ...boundingBox };
       documentRect.x += scrollOffset.x;
       documentRect.y += scrollOffset.y;
-      return await this._screenshot(format, helper.enclosingIntRect(documentRect), undefined, options, overridenViewportSize, originalViewportSize);
+      const buffer = await this._screenshot(progress, format, helper.enclosingIntRect(documentRect), undefined, options);
+      progress.throwIfAborted(); // Avoid restoring after failure - should be done by cleanup.
+      if (overridenViewportSize)
+        await this._restoreViewport(originalViewportSize);
+      return buffer;
     }).catch(rewriteError);
   }
 
-  private async _screenshot(format: 'png' | 'jpeg', documentRect: types.Rect | undefined, viewportRect: types.Rect | undefined, options: types.ElementScreenshotOptions, overridenViewportSize: types.Size | null, originalViewportSize: types.Size | null): Promise<Buffer> {
+  private async _screenshot(progress: Progress, format: 'png' | 'jpeg', documentRect: types.Rect | undefined, viewportRect: types.Rect | undefined, options: types.ElementScreenshotOptions): Promise<Buffer> {
+    if ((options as any).__testHookBeforeScreenshot)
+      await (options as any).__testHookBeforeScreenshot();
+    progress.throwIfAborted(); // Screenshotting is expensive - avoid extra work.
     const shouldSetDefaultBackground = options.omitBackground && format === 'png';
-    if (shouldSetDefaultBackground)
+    if (shouldSetDefaultBackground) {
       await this._page._delegate.setBackgroundColor({ r: 0, g: 0, b: 0, a: 0});
+      progress.cleanupWhenAborted(() => this._page._delegate.setBackgroundColor());
+    }
     const buffer = await this._page._delegate.takeScreenshot(format, documentRect, viewportRect, options.quality);
+    progress.throwIfAborted(); // Avoid restoring after failure - should be done by cleanup.
     if (shouldSetDefaultBackground)
       await this._page._delegate.setBackgroundColor();
-    if (overridenViewportSize) {
+    progress.throwIfAborted(); // Avoid side effects.
-      assert(!this._page._delegate.canScreenshotOutsideViewport());
-      if (originalViewportSize)
-        await this._page.setViewportSize(originalViewportSize);
-      else
-        await this._page._delegate.resetViewport();
-    }
     if (options.path)
       await util.promisify(fs.writeFile)(options.path, buffer);
+    if ((options as any).__testHookAfterScreenshot)
+      await (options as any).__testHookAfterScreenshot();
     return buffer;
   }
+
+  private async _restoreViewport(originalViewportSize: types.Size | null) {
+    assert(!this._page._delegate.canScreenshotOutsideViewport());
+    if (originalViewportSize)
+      await this._page.setViewportSize(originalViewportSize);
+    else
+      await this._page._delegate.resetViewport();
+  }
 }
 
 class TaskQueue {

src/types.ts

@@ -48,7 +48,7 @@ export type WaitForNavigationOptions = TimeoutOptions & {
   url?: URLMatch
 };
 
-export type ElementScreenshotOptions = {
+export type ElementScreenshotOptions = TimeoutOptions & {
   type?: 'png' | 'jpeg',
   path?: string,
   quality?: number,

test/browsercontext.spec.js

@@ -85,10 +85,7 @@ describe('BrowserContext', function() {
   it('should propagate default viewport to the page', async({ browser }) => {
     const context = await browser.newContext({ viewport: { width: 456, height: 789 } });
     const page = await context.newPage();
-    expect(page.viewportSize().width).toBe(456);
+    await utils.verifyViewport(page, 456, 789);
-    expect(page.viewportSize().height).toBe(789);
-    expect(await page.evaluate('window.innerWidth')).toBe(456);
-    expect(await page.evaluate('window.innerHeight')).toBe(789);
     await context.close();
   });
   it('should make a copy of default viewport', async({ browser }) => {
@@ -96,10 +93,7 @@ describe('BrowserContext', function() {
     const context = await browser.newContext({ viewport });
     viewport.width = 567;
     const page = await context.newPage();
-    expect(page.viewportSize().width).toBe(456);
+    await utils.verifyViewport(page, 456, 789);
-    expect(page.viewportSize().height).toBe(789);
-    expect(await page.evaluate('window.innerWidth')).toBe(456);
-    expect(await page.evaluate('window.innerHeight')).toBe(789);
     await context.close();
   });
   it('should respect deviceScaleFactor', async({ browser }) => {

test/defaultbrowsercontext.spec.js

@@ -135,15 +135,9 @@ describe('launchPersistentContext()', function() {
   });
   it('should support viewport option', async state => {
     let { page, context } = await launch(state, {viewport: { width: 456, height: 789 }});
-    expect(page.viewportSize().width).toBe(456);
+    await utils.verifyViewport(page, 456, 789);
-    expect(page.viewportSize().height).toBe(789);
-    expect(await page.evaluate('window.innerWidth')).toBe(456);
-    expect(await page.evaluate('window.innerHeight')).toBe(789);
     page = await context.newPage();
-    expect(page.viewportSize().width).toBe(456);
+    await utils.verifyViewport(page, 456, 789);
-    expect(page.viewportSize().height).toBe(789);
-    expect(await page.evaluate('window.innerWidth')).toBe(456);
-    expect(await page.evaluate('window.innerHeight')).toBe(789);
     await close(state);
   });
   it('should support deviceScaleFactor option', async state => {

test/emulation.spec.js

@@ -21,17 +21,13 @@ const iPhone = playwright.devices['iPhone 6'];
 const iPhoneLandscape = playwright.devices['iPhone 6 landscape'];
 
 describe('BrowserContext({viewport})', function() {
-  it('should get the proper viewport size', async({page, server}) => {
+  it('should get the proper default viewport size', async({page, server}) => {
-    expect(page.viewportSize()).toEqual({width: 1280, height: 720});
+    await utils.verifyViewport(page, 1280, 720);
-    expect(await page.evaluate(() => window.innerWidth)).toBe(1280);
-    expect(await page.evaluate(() => window.innerHeight)).toBe(720);
   });
   it('should set the proper viewport size', async({page, server}) => {
-    expect(page.viewportSize()).toEqual({width: 1280, height: 720});
+    await utils.verifyViewport(page, 1280, 720);
     await page.setViewportSize({width: 123, height: 456});
-    expect(page.viewportSize()).toEqual({width: 123, height: 456});
+    await utils.verifyViewport(page, 123, 456);
-    expect(await page.evaluate(() => window.innerWidth)).toBe(123);
-    expect(await page.evaluate(() => window.innerHeight)).toBe(456);
   });
   it('should emulate device width', async({page, server}) => {
     expect(page.viewportSize()).toEqual({width: 1280, height: 720});

test/screenshot.spec.js

@@ -15,7 +15,8 @@
  * limitations under the License.
  */
 
-const {FFOX, CHROMIUM, WEBKIT} = require('./utils').testOptions(browserType);
+const utils = require('./utils');
+const {FFOX, CHROMIUM, WEBKIT} = utils.testOptions(browserType);
 const {PNG} = require('pngjs');
 
 // Firefox headful produces a different image.
@@ -80,7 +81,7 @@ describe.skip(ffheadful)('Page.screenshot', function() {
         height: 100
       }
     }).catch(error => error);
-    expect(screenshotError.message).toBe('Clipped area is either empty or outside the resulting image');
+    expect(screenshotError.message).toContain('Clipped area is either empty or outside the resulting image');
   });
   it('should run in parallel', async({page, server, golden}) => {
     await page.setViewportSize({width: 500, height: 500});
@@ -112,8 +113,7 @@ describe.skip(ffheadful)('Page.screenshot', function() {
     await page.goto(server.PREFIX + '/grid.html');
     const screenshot = await page.screenshot({ fullPage: true });
     expect(screenshot).toBeInstanceOf(Buffer);
-    expect(page.viewportSize().width).toBe(500);
+    await utils.verifyViewport(page, 500, 500);
-    expect(page.viewportSize().height).toBe(500);
   });
   it('should run in parallel in multiple pages', async({page, server, context, golden}) => {
     const N = 5;
@@ -281,7 +281,7 @@ describe.skip(ffheadful)('ElementHandle.screenshot', function() {
     const screenshots = await Promise.all(promises);
     expect(screenshots[2]).toBeGolden(golden('screenshot-element-larger-than-viewport.png'));
 
-    expect(await page.evaluate(() => ({ w: window.innerWidth, h: window.innerHeight }))).toEqual({ w: 500, h: 500 });
+    await utils.verifyViewport(page, 500, 500);
   });
   it('should capture full element when larger than viewport', async({page, golden}) => {
     await page.setViewportSize({width: 500, height: 500});
@@ -307,7 +307,7 @@ describe.skip(ffheadful)('ElementHandle.screenshot', function() {
     const screenshot = await elementHandle.screenshot();
     expect(screenshot).toBeGolden(golden('screenshot-element-larger-than-viewport.png'));
 
-    expect(await page.evaluate(() => ({ w: window.innerWidth, h: window.innerHeight }))).toEqual({ w: 500, h: 500 });
+    await utils.verifyViewport(page, 500, 500);
   });
   it('should scroll element into view', async({page, golden}) => {
     await page.setViewportSize({width: 500, height: 500});
@@ -375,11 +375,30 @@ describe.skip(ffheadful)('ElementHandle.screenshot', function() {
     const screenshotError = await elementHandle.screenshot().catch(error => error);
     expect(screenshotError.message).toContain('Element is not attached to the DOM');
   });
-  it('should not hang with zero width/height element', async({page, server}) => {
+  it('should timeout waiting for visible', async({page, server}) => {
     await page.setContent('<div style="width: 50px; height: 0"></div>');
     const div = await page.$('div');
-    const error = await div.screenshot().catch(e => e);
+    const error = await div.screenshot({ timeout: 3000 }).catch(e => e);
-    expect(error.message).toBe('Node has 0 height.');
+    expect(error.message).toContain('Timeout 3000ms exceeded during elementHandle.screenshot');
+    expect(error.message).toContain('element is not visible');
+  });
+  it('should wait for visible', async({page, server, golden}) => {
+    await page.setViewportSize({width: 500, height: 500});
+    await page.goto(server.PREFIX + '/grid.html');
+    await page.evaluate(() => window.scrollBy(50, 100));
+    const elementHandle = await page.$('.box:nth-of-type(3)');
+    await elementHandle.evaluate(e => e.style.visibility = 'hidden');
+    let done = false;
+    const promise = elementHandle.screenshot().then(buffer => {
+      done = true;
+      return buffer;
+    });
+    for (let i = 0; i < 10; i++)
+      await page.evaluate(() => new Promise(f => requestAnimationFrame(f)));
+    expect(done).toBe(false);
+    await elementHandle.evaluate(e => e.style.visibility = 'visible');
+    const screenshot = await promise;
+    expect(screenshot).toBeGolden(golden('screenshot-element-bounding-box.png'));
   });
   it('should work for an element with fractional dimensions', async({page, golden}) => {
     await page.setContent('<div style="width:48.51px;height:19.8px;border:1px solid black;"></div>');
@@ -449,16 +468,33 @@ describe.skip(ffheadful)('ElementHandle.screenshot', function() {
   it('should restore default viewport after fullPage screenshot', async({ browser }) => {
     const context = await browser.newContext({ viewport: { width: 456, height: 789 } });
     const page = await context.newPage();
-    expect(page.viewportSize().width).toBe(456);
+    await utils.verifyViewport(page, 456, 789);
-    expect(page.viewportSize().height).toBe(789);
-    expect(await page.evaluate('window.innerWidth')).toBe(456);
-    expect(await page.evaluate('window.innerHeight')).toBe(789);
     const screenshot = await page.screenshot({ fullPage: true });
     expect(screenshot).toBeInstanceOf(Buffer);
-    expect(page.viewportSize().width).toBe(456);
+    await utils.verifyViewport(page, 456, 789);
-    expect(page.viewportSize().height).toBe(789);
+    await context.close();
-    expect(await page.evaluate('window.innerWidth')).toBe(456);
+  });
-    expect(await page.evaluate('window.innerHeight')).toBe(789);
+  it('should restore viewport after page screenshot and exception', async({ browser, server }) => {
+    const context = await browser.newContext({ viewport: { width: 350, height: 360 } });
+    const page = await context.newPage();
+    await page.goto(server.PREFIX + '/grid.html');
+    const __testHookBeforeScreenshot = () => { throw new Error('oh my') };
+    const error = await page.screenshot({ fullPage: true, __testHookBeforeScreenshot }).catch(e => e);
+    expect(error.message).toContain('oh my');
+    await utils.verifyViewport(page, 350, 360);
+    await context.close();
+  });
+  it('should restore viewport after page screenshot and timeout', async({ browser, server }) => {
+    const context = await browser.newContext({ viewport: { width: 350, height: 360 } });
+    const page = await context.newPage();
+    await page.goto(server.PREFIX + '/grid.html');
+    const __testHookAfterScreenshot = () => new Promise(f => setTimeout(f, 5000));
+    const error = await page.screenshot({ fullPage: true, __testHookAfterScreenshot, timeout: 3000 }).catch(e => e);
+    expect(error.message).toContain('Timeout 3000ms exceeded during page.screenshot');
+    await utils.verifyViewport(page, 350, 360);
+    await page.setViewportSize({ width: 400, height: 400 });
+    await page.waitForTimeout(3000); // Give it some time to wrongly restore previous viewport.
+    await utils.verifyViewport(page, 400, 400);
     await context.close();
   });
   it('should take element screenshot when default viewport is null and restore back', async({server, browser}) => {
@@ -490,4 +526,15 @@ describe.skip(ffheadful)('ElementHandle.screenshot', function() {
     expect(sizeBefore.height).toBe(sizeAfter.height);
     await context.close();
   });
+  it('should restore viewport after element screenshot and exception', async({server, browser}) => {
+    const context = await browser.newContext({ viewport: { width: 350, height: 360 } });
+    const page = await context.newPage();
+    await page.setContent(`<div style="width:600px;height:600px;"></div>`);
+    const elementHandle = await page.$('div');
+    const __testHookBeforeScreenshot = () => { throw new Error('oh my') };
+    const error = await elementHandle.screenshot({ __testHookBeforeScreenshot }).catch(e => e);
+    expect(error.message).toContain('oh my');
+    await utils.verifyViewport(page, 350, 360);
+    await context.close();
+  });
 });

test/utils.js

@@ -87,6 +87,13 @@ const utils = module.exports = {
     return result;
   },
 
+  verifyViewport: async (page, width, height) => {
+    expect(page.viewportSize().width).toBe(width);
+    expect(page.viewportSize().height).toBe(height);
+    expect(await page.evaluate('window.innerWidth')).toBe(width);
+    expect(await page.evaluate('window.innerHeight')).toBe(height);
+  },
+
   initializeFlakinessDashboardIfNeeded: async function(testRunner) {
     // Generate testIDs for all tests and verify they don't clash.
     // This will add |test.testId| for every test.