From 021f51cc1fb0aed5b33d998cd91a630715d1672b Mon Sep 17 00:00:00 2001
From: Pavel Feldman <pavel.feldman@gmail.com>
Date: Tue, 8 Jun 2021 11:22:07 -0700
Subject: [PATCH] feat(trace): add a trace option into the test runner (#6961)

Co-authored-by: Andrey Lushnikov <aslushnikov@gmail.com>
---
 docs/src/test-configuration.md |  5 +++++
 src/test/index.ts              | 19 ++++++++++++++++++-
 types/test.d.ts                |  9 +++++++++
 3 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/docs/src/test-configuration.md b/docs/src/test-configuration.md
index 18209da7a6..1548d6a606 100644
--- a/docs/src/test-configuration.md
+++ b/docs/src/test-configuration.md
@@ -17,6 +17,11 @@ You can specify any options either locally in a test file, or globally in the co
   - `'off'` - Do not capture screenshots.
   - `'on'` - Capture screenshot after each test.
   - `'only-on-failure'` - Capture screenshot after each test failure.
+- `trace` option - whether to record trace for each test, off by default. Trace will appear in the test output directory, typically `test-results`.
+  - `'off'` - Do not record trace.
+  - `'on'` - Record trace for each test.
+  - `'retain-on-failure'` - Record trace for each test, but remove it from successful test runs.
+  - `'retry-with-trace'` - Record trace only when retrying a test.
 - `video` option - whether to record video for each test, off by default. Video will appear in the test output directory, typically `test-results`.
   - `'off'` - Do not record video.
   - `'on'` - Record video for each test.
diff --git a/src/test/index.ts b/src/test/index.ts
index 7e1a0839bf..dec523ecc2 100644
--- a/src/test/index.ts
+++ b/src/test/index.ts
@@ -15,6 +15,7 @@
  */
 
 import * as fs from 'fs';
+import * as path from 'path';
 import type { LaunchOptions, BrowserContextOptions, Page } from '../../types/types';
 import type { TestType, PlaywrightTestArgs, PlaywrightTestOptions, PlaywrightWorkerArgs, PlaywrightWorkerOptions } from '../../types/test';
 import { rootTestType } from './testType';
@@ -47,6 +48,7 @@ export const test = _baseTest.extend<PlaywrightTestArgs & PlaywrightTestOptions,
 
   screenshot: 'off',
   video: 'off',
+  trace: 'off',
   acceptDownloads: undefined,
   bypassCSP: undefined,
   colorScheme: undefined,
@@ -68,7 +70,7 @@ export const test = _baseTest.extend<PlaywrightTestArgs & PlaywrightTestOptions,
   viewport: undefined,
   contextOptions: {},
 
-  context: async ({ browser, screenshot, video, acceptDownloads, bypassCSP, colorScheme, deviceScaleFactor, extraHTTPHeaders, hasTouch, geolocation, httpCredentials, ignoreHTTPSErrors, isMobile, javaScriptEnabled, locale, offline, permissions, proxy, storageState, viewport, timezoneId, userAgent, contextOptions }, use, testInfo) => {
+  context: async ({ browser, screenshot, trace, video, acceptDownloads, bypassCSP, colorScheme, deviceScaleFactor, extraHTTPHeaders, hasTouch, geolocation, httpCredentials, ignoreHTTPSErrors, isMobile, javaScriptEnabled, locale, offline, permissions, proxy, storageState, viewport, timezoneId, userAgent, contextOptions }, use, testInfo) => {
     testInfo.snapshotSuffix = process.platform;
     if (process.env.PWDEBUG)
       testInfo.setTimeout(0);
@@ -122,9 +124,24 @@ export const test = _baseTest.extend<PlaywrightTestArgs & PlaywrightTestOptions,
     const allPages: Page[] = [];
     context.on('page', page => allPages.push(page));
 
+    const collectingTrace = trace === 'on' || trace === 'retain-on-failure' || (trace === 'retry-with-trace' && testInfo.retry);
+    if (collectingTrace) {
+      const name = path.relative(testInfo.project.outputDir, testInfo.outputDir).replace(/[\/\\]/g, '-');
+      await context.tracing.start({ name, screenshots: true, snapshots: true });
+    }
+
     await use(context);
 
     const testFailed = testInfo.status !== testInfo.expectedStatus;
+
+    const saveTrace = trace === 'on' || (testFailed && trace === 'retain-on-failure') || (trace === 'retry-with-trace' && testInfo.retry);
+    if (saveTrace) {
+      const tracePath = testInfo.outputPath(`trace.zip`);
+      await context.tracing.stop({ path: tracePath });
+    } else if (collectingTrace) {
+      await context.tracing.stop();
+    }
+
     if (screenshot === 'on' || (screenshot === 'only-on-failure' && testFailed)) {
       await Promise.all(allPages.map((page, index) => {
         const screenshotPath = testInfo.outputPath(`test-${testFailed ? 'failed' : 'finished'}-${++index}.png`);
diff --git a/types/test.d.ts b/types/test.d.ts
index f3756d3df7..47b1276942 100644
--- a/types/test.d.ts
+++ b/types/test.d.ts
@@ -961,6 +961,15 @@ export type PlaywrightTestOptions = {
    */
   screenshot: 'off' | 'on' | 'only-on-failure';
 
+  /**
+   * Whether to record trace for each test, off by default.
+   * - `off`: Do not record trace.
+   * - `on`: Record trace for each test.
+   * - `retain-on-failure`: Record trace for each test, but remove trace from successful test run.
+   * - `retry-with-trace`: Record trace only when retrying a test.
+   */
+  trace: 'off' | 'on' | 'retain-on-failure' | 'retry-with-trace';
+
   /**
   * Whether to record video for each test, off by default.
   * - `off`: Do not record video.