From 220d79af1e0539b3dcdcce0d57ac3458f1c0ed60 Mon Sep 17 00:00:00 2001
From: Dmitry Gozman <dgozman@gmail.com>
Date: Mon, 15 Jul 2024 11:27:12 +0100
Subject: [PATCH] fix(tracing): move tracing scope to only cover test

This makes tracing not cover:
- worker fixtures setup and teardown
- beforeAll and afterAll hooks
---
 packages/playwright/src/worker/workerMain.ts  | 10 ++++--
 .../playwright-test/playwright.trace.spec.ts  | 33 +++++++++++++++++++
 2 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/packages/playwright/src/worker/workerMain.ts b/packages/playwright/src/worker/workerMain.ts
index 2f0dcc424d..2b2a42e1af 100644
--- a/packages/playwright/src/worker/workerMain.ts
+++ b/packages/playwright/src/worker/workerMain.ts
@@ -319,7 +319,6 @@ export class WorkerMain extends ProcessRunner {
         if (typeof traceFixtureRegistration.fn === 'function')
           throw new Error(`"trace" option cannot be a function`);
         await testInfo._tracing.startIfNeeded(traceFixtureRegistration.fn);
-        await testLifecycleInstrumentation()?.onTestBegin?.();
       });
 
       if (this._isStopped || isSkipped) {
@@ -339,6 +338,8 @@ export class WorkerMain extends ProcessRunner {
         for (const suite of suites)
           await this._runBeforeAllHooksForSuite(suite, testInfo);
 
+        await testInfo._runAsStage({ title: 'start tracing', runnable: { type: 'test', slot: tracingSlot } }, async () => testLifecycleInstrumentation()?.onTestBegin?.());
+
         // Run "beforeEach" hooks. Once started with "beforeEach", we must run all "afterEach" hooks as well.
         shouldRunAfterEachHooks = true;
         await this._runEachHooksForSuites(suites, 'beforeEach', testInfo);
@@ -403,6 +404,12 @@ export class WorkerMain extends ProcessRunner {
         firstAfterHooksError = firstAfterHooksError ?? error;
       }
 
+      try {
+        await testInfo._runAsStage({ title: 'stop tracing', runnable: { type: 'test', slot: tracingSlot } }, async () => testLifecycleInstrumentation()?.onTestEnd?.());
+      } catch (error) {
+        firstAfterHooksError = firstAfterHooksError ?? error;
+      }
+
       // Run "afterAll" hooks for suites that are not shared with the next test.
       // In case of failure the worker will be stopped and we have to make sure that afterAll
       // hooks run before worker fixtures teardown.
@@ -461,7 +468,6 @@ export class WorkerMain extends ProcessRunner {
     }
 
     await testInfo._runAsStage({ title: 'stop tracing', runnable: { type: 'test', slot: tracingSlot } }, async () => {
-      await testLifecycleInstrumentation()?.onTestEnd?.();
       await testInfo._tracing.stopIfNeeded();
     }).catch(() => {});  // Ignore the top-level error, it is already inside TestInfo.errors.
 
diff --git a/tests/playwright-test/playwright.trace.spec.ts b/tests/playwright-test/playwright.trace.spec.ts
index 01d4081af7..e6b32080c3 100644
--- a/tests/playwright-test/playwright.trace.spec.ts
+++ b/tests/playwright-test/playwright.trace.spec.ts
@@ -1268,3 +1268,36 @@ test('should take a screenshot-on-failure in workerStorageState', async ({ runIn
   expect(result.failed).toBe(1);
   expect(fs.existsSync(test.info().outputPath('test-results', 'a-fail', 'test-failed-1.png'))).toBeTruthy();
 });
+
+test('should record trace in a manually create context in a failed test', async ({ runInlineTest }) => {
+  test.info().annotations.push({ type: 'issue', description: 'https://github.com/microsoft/playwright/issues/31541' });
+
+  const result = await runInlineTest({
+    'a.spec.ts': `
+      import { test, expect } from '@playwright/test';
+      test('fail', async ({ browser }) => {
+        const page = await browser.newPage();
+        await page.setContent('<script>console.log("from the page");</script>');
+        expect(1).toBe(2);
+      });
+    `,
+  }, { trace: 'on' });
+  expect(result.exitCode).toBe(1);
+  expect(result.failed).toBe(1);
+
+  const tracePath = test.info().outputPath('test-results', 'a-fail', 'trace.zip');
+  const trace = await parseTrace(tracePath);
+  expect(trace.actionTree).toEqual([
+    'Before Hooks',
+    '  fixture: browser',
+    '    browserType.launch',
+    'browser.newPage',
+    'page.setContent',
+    'expect.toBe',
+    'After Hooks',
+    'Worker Cleanup',
+    '  fixture: browser',
+  ]);
+  // Check console events to make sure that library trace is recorded.
+  expect(trace.events).toContainEqual(expect.objectContaining({ type: 'console', text: 'from the page' }));
+});