diff --git a/packages/playwright-core/src/server/screenshotter.ts b/packages/playwright-core/src/server/screenshotter.ts index 4e82cfe38cba3..c544069445b72 100644 --- a/packages/playwright-core/src/server/screenshotter.ts +++ b/packages/playwright-core/src/server/screenshotter.ts @@ -222,11 +222,18 @@ export class Screenshotter { })); } - async _maskElements(progress: Progress, options: ScreenshotOptions) { + async _maskElements(progress: Progress, options: ScreenshotOptions): Promise<() => Promise> { + const framesToParsedSelectors: MultiMap = new MultiMap(); + + const cleanup = async () => { + await Promise.all([...framesToParsedSelectors.keys()].map(async frame => { + await frame.hideHighlight(); + })); + }; + if (!options.mask || !options.mask.length) - return false; + return cleanup; - const framesToParsedSelectors: MultiMap = new MultiMap(); await Promise.all((options.mask || []).map(async ({ frame, selector }) => { const pair = await frame.resolveFrameForSelectorNoWait(selector); if (pair) @@ -237,8 +244,8 @@ export class Screenshotter { await Promise.all([...framesToParsedSelectors.keys()].map(async frame => { await frame.maskSelectors(framesToParsedSelectors.get(frame)); })); - progress.cleanupWhenAborted(() => this._page.hideHighlight()); - return true; + progress.cleanupWhenAborted(cleanup); + return cleanup; } private async _screenshot(progress: Progress, format: 'png' | 'jpeg', documentRect: types.Rect | undefined, viewportRect: types.Rect | undefined, fitsViewport: boolean | undefined, options: ScreenshotOptions): Promise { @@ -252,14 +259,13 @@ export class Screenshotter { } progress.throwIfAborted(); // Avoid extra work. - const hasHighlight = await this._maskElements(progress, options); + const cleanupHighlight = await this._maskElements(progress, options); progress.throwIfAborted(); // Avoid extra work. const buffer = await this._page._delegate.takeScreenshot(progress, format, documentRect, viewportRect, options.quality, fitsViewport); progress.throwIfAborted(); // Avoid restoring after failure - should be done by cleanup. - if (hasHighlight) - await this._page.hideHighlight(); + await cleanupHighlight(); progress.throwIfAborted(); // Avoid restoring after failure - should be done by cleanup. if (shouldSetDefaultBackground) diff --git a/tests/page/page-screenshot.spec.ts b/tests/page/page-screenshot.spec.ts index 1247c28383aca..ff36b98cf0a13 100644 --- a/tests/page/page-screenshot.spec.ts +++ b/tests/page/page-screenshot.spec.ts @@ -444,6 +444,19 @@ it.describe('page screenshot', () => { await route.fulfill({ body: '' }); await done; }); + + it('should work when subframe used document.open after a weird url', async ({ page, server }) => { + await page.goto(server.EMPTY_PAGE); + await page.evaluate(() => { + const iframe = document.createElement('iframe'); + iframe.src = 'javascript:hi'; + document.body.appendChild(iframe); + iframe.contentDocument.open(); + iframe.contentDocument.write('Hello'); + iframe.contentDocument.close(); + }); + await page.screenshot({ mask: [ page.locator('non-existent') ] }); + }); }); });