fix(test runner): screenshot immediately after failure (#15159)

Previously, screenshot was taken after hooks and fixtures teardown.
However, hooks can easily modify the state of the page, and
screenshot would not reflect the moment of failure.

Instead, we take screenshots immediately after the test function
finishes with an error.

packages/playwright-test/src/index.ts

@@ -25,6 +25,7 @@ export { expect } from './expect';
 export const _baseTest: TestType<{}, {}> = rootTestType.test;
 export { addRunnerPlugin as _addRunnerPlugin } from './plugins';
 import * as outOfProcess from 'playwright-core/lib/outofprocess';
+import type { TestInfoImpl } from './testInfo';
 
 if ((process as any)['__pw_initiator__']) {
   const originalStackTraceLimit = Error.stackTraceLimit;
@@ -249,6 +250,7 @@ export const test = _baseTest.extend<TestFixtures, WorkerFixtures>({
     const temporaryTraceFiles: string[] = [];
     const temporaryScreenshots: string[] = [];
     const createdContexts = new Set<BrowserContext>();
+    const testInfoImpl = testInfo as TestInfoImpl;
 
     const createInstrumentationListener = (context?: BrowserContext) => {
       return {
@@ -260,9 +262,8 @@ export const test = _baseTest.extend<TestFixtures, WorkerFixtures>({
             context?.setDefaultNavigationTimeout(0);
             context?.setDefaultTimeout(0);
           }
-          const testInfoImpl = testInfo as any;
           const step = testInfoImpl._addStep({
-            location: stackTrace?.frames[0],
+            location: stackTrace?.frames[0] as any,
             category: 'pw:api',
             title: apiCall,
             canHaveChildren: false,
@@ -320,16 +321,29 @@ export const test = _baseTest.extend<TestFixtures, WorkerFixtures>({
       }
     };
 
+    const screenshottedSymbol = Symbol('screenshotted');
+    const screenshotPage = async (page: Page) => {
+      if ((page as any)[screenshottedSymbol])
+        return;
+      (page as any)[screenshottedSymbol] = true;
+      const screenshotPath = path.join(_artifactsDir(), createGuid() + '.png');
+      temporaryScreenshots.push(screenshotPath);
+      await page.screenshot({ timeout: 5000, path: screenshotPath }).catch(() => {});
+    };
+
+    const screenshotOnTestFailure = async () => {
+      const contexts: BrowserContext[] = [];
+      for (const browserType of [playwright.chromium, playwright.firefox, playwright.webkit])
+        contexts.push(...(browserType as any)._contexts);
+      await Promise.all(contexts.map(ctx => Promise.all(ctx.pages().map(screenshotPage))));
+    };
+
     const onWillCloseContext = async (context: BrowserContext) => {
       await stopTracing(context.tracing);
       if (screenshot === 'on' || screenshot === 'only-on-failure') {
         // Capture screenshot for now. We'll know whether we have to preserve them
         // after the test finishes.
-        await Promise.all(context.pages().map(async page => {
-          const screenshotPath = path.join(_artifactsDir(), createGuid() + '.png');
-          temporaryScreenshots.push(screenshotPath);
-          await page.screenshot({ timeout: 5000, path: screenshotPath }).catch(() => {});
-        }));
+        await Promise.all(context.pages().map(screenshotPage));
       }
     };
 
@@ -352,6 +366,8 @@ export const test = _baseTest.extend<TestFixtures, WorkerFixtures>({
       const existingApiRequests: APIRequestContext[] =  Array.from((playwright.request as any)._contexts as Set<APIRequestContext>);
       await Promise.all(existingApiRequests.map(onDidCreateRequestContext));
     }
+    if (screenshot === 'on' || screenshot === 'only-on-failure')
+      testInfoImpl._onTestFailureImmediateCallbacks.set(screenshotOnTestFailure, 'Screenshot on failure');
 
     // 2. Run the test.
     await use();
@@ -391,6 +407,7 @@ export const test = _baseTest.extend<TestFixtures, WorkerFixtures>({
     const leftoverApiRequests: APIRequestContext[] =  Array.from((playwright.request as any)._contexts as Set<APIRequestContext>);
     (playwright.request as any)._onDidCreateContext = undefined;
     (playwright.request as any)._onWillCloseContext = undefined;
+    testInfoImpl._onTestFailureImmediateCallbacks.delete(screenshotOnTestFailure);
 
     const stopTraceChunk = async (tracing: Tracing): Promise<boolean> => {
       // When we timeout during context.close(), we might end up with context still alive
@@ -409,8 +426,13 @@ export const test = _baseTest.extend<TestFixtures, WorkerFixtures>({
     await Promise.all(leftoverContexts.map(async context => {
       if (!await stopTraceChunk(context.tracing))
         return;
-      if (captureScreenshots)
-        await Promise.all(context.pages().map(page => page.screenshot({ timeout: 5000, path: addScreenshotAttachment() }).catch(() => {})));
+      if (captureScreenshots) {
+        await Promise.all(context.pages().map(async page => {
+          if ((page as any)[screenshottedSymbol])
+            return;
+          await page.screenshot({ timeout: 5000, path: addScreenshotAttachment() }).catch(() => {});
+        }));
+      }
     }).concat(leftoverApiRequests.map(async context => {
       const tracing = (context as any)._tracing as Tracing;
       await stopTraceChunk(tracing);

packages/playwright-test/src/testInfo.ts

@@ -33,6 +33,7 @@ export class TestInfoImpl implements TestInfo {
   readonly _startWallTime: number;
   private _hasHardError: boolean = false;
   readonly _screenshotsDir: string;
+  readonly _onTestFailureImmediateCallbacks = new Map<() => Promise<void>, string>(); // fn -> title
 
   // ------------ TestInfo fields ------------
   readonly repeatEachIndex: number;
@@ -224,6 +225,10 @@ export class TestInfoImpl implements TestInfo {
     }
   }
 
+  _isFailure() {
+    return this.status !== 'skipped' && this.status !== this.expectedStatus;
+  }
+
   // ------------ TestInfo methods ------------
 
   async attach(name: string, options: { path?: string, body?: string | Buffer, contentType?: string } = {}) {

packages/playwright-test/src/workerRunner.ts

@@ -385,6 +385,22 @@ export class WorkerRunner extends EventEmitter {
       // Note: do not wrap all teardown steps together, because failure in any of them
       // does not prevent further teardown steps from running.
 
+      // Run "immediately upon test failure" callbacks.
+      if (testInfo._isFailure()) {
+        const onFailureError = await testInfo._runFn(async () => {
+          testInfo._timeoutManager.setCurrentRunnable({ type: 'test', slot: afterHooksSlot });
+          for (const [fn, title] of testInfo._onTestFailureImmediateCallbacks) {
+            await testInfo._runAsStep(fn, {
+              category: 'hook',
+              title,
+              canHaveChildren: true,
+              forceNoParent: false,
+            });
+          }
+        });
+        firstAfterHooksError = firstAfterHooksError || onFailureError;
+      }
+
       // Run "afterEach" hooks, unless we failed at beforeAll stage.
       if (shouldRunAfterEachHooks) {
         const afterEachError = await testInfo._runFn(() => this._runEachHooksForSuites(reversedSuites, 'afterEach', testInfo, afterHooksSlot));
@@ -395,9 +411,8 @@ export class WorkerRunner extends EventEmitter {
       const nextSuites = new Set(getSuites(nextTest));
       // In case of failure the worker will be stopped and we have to make sure that afterAll
       // hooks run before test fixtures teardown.
-      const isFailure = testInfo.status !== 'skipped' && testInfo.status !== testInfo.expectedStatus;
       for (const suite of reversedSuites) {
-        if (!nextSuites.has(suite) || isFailure) {
+        if (!nextSuites.has(suite) || testInfo._isFailure()) {
           const afterAllError = await this._runAfterAllHooksForSuite(suite, testInfo);
           firstAfterHooksError = firstAfterHooksError || afterAllError;
         }
@@ -409,8 +424,7 @@ export class WorkerRunner extends EventEmitter {
       firstAfterHooksError = firstAfterHooksError || testScopeError;
     });
 
-    const isFailure = testInfo.status !== 'skipped' && testInfo.status !== testInfo.expectedStatus;
-    if (isFailure)
+    if (testInfo._isFailure())
       this._isStopped = true;
 
     if (this._isStopped) {
@@ -439,7 +453,7 @@ export class WorkerRunner extends EventEmitter {
     this.emit('testEnd', buildTestEndPayload(testInfo));
 
     const preserveOutput = this._loader.fullConfig().preserveOutput === 'always' ||
-      (this._loader.fullConfig().preserveOutput === 'failures-only' && isFailure);
+      (this._loader.fullConfig().preserveOutput === 'failures-only' && testInfo._isFailure());
     if (!preserveOutput)
       await removeFolderAsync(testInfo.outputDir).catch(e => {});
   }

tests/playwright-test/playwright.artifacts.spec.ts

@@ -278,3 +278,26 @@ test('should work with trace: on-first-retry', async ({ runInlineTest }, testInf
     'report.json',
   ]);
 });
+
+test('should take screenshot when page is closed in afterEach', async ({ runInlineTest }, testInfo) => {
+  const result = await runInlineTest({
+    'playwright.config.ts': `
+      module.exports = { use: { screenshot: 'on' } };
+    `,
+    'a.spec.ts': `
+      const { test } = pwt;
+
+      test.afterEach(async ({ page }) => {
+        await page.close();
+      });
+
+      test('fails', async ({ page }) => {
+        expect(1).toBe(2);
+      });
+    `,
+  }, { workers: 1 });
+
+  expect(result.exitCode).toBe(1);
+  expect(result.failed).toBe(1);
+  expect(fs.existsSync(testInfo.outputPath('test-results', 'a-fails', 'test-failed-1.png'))).toBeTruthy();
+});