feat(ui-mode): update individual snapshots

packages/playwright/src/isomorphic/testServerConnection.ts

@@ -219,6 +219,10 @@ export class TestServerConnection implements TestServerInterface, TestServerInte
     return await this._sendMessage('clearCache', params);
   }
 
+  async updateSnapshot(params: Parameters<TestServerInterface['updateSnapshot']>[0]): ReturnType<TestServerInterface['updateSnapshot']> {
+    return await this._sendMessage('updateSnapshot', params);
+  }
+
   async listFiles(params: Parameters<TestServerInterface['listFiles']>[0]): ReturnType<TestServerInterface['listFiles']> {
     return await this._sendMessage('listFiles', params);
   }

packages/playwright/src/isomorphic/testServerInterface.ts

@@ -22,6 +22,19 @@ import type * as reporterTypes from '../../types/testReporter';
 
 export type ReportEntry = JsonEvent;
 
+export type UpdateSnapshotParams = {
+  testId: string;
+  resultId: string;
+  actual: {
+    name: string;
+    contentType: string;
+  };
+  expected: {
+    name: string;
+    contentType: string;
+  };
+};
+
 export interface TestServerInterface {
   initialize(params: {
     serializer?: string,
@@ -57,6 +70,8 @@ export interface TestServerInterface {
 
   clearCache(params: {}): Promise<void>;
 
+  updateSnapshot(params: UpdateSnapshotParams): Promise<void>;
+
   listFiles(params: {
     projects?: string[];
   }): Promise<{

packages/playwright/src/runner/testServer.ts

@@ -14,13 +14,15 @@
  * limitations under the License.
  */
 
-import util from 'util';
 import debug from 'debug';
+import fs from 'fs';
 import open from 'open';
+import util from 'util';
 
 import { server as coreServer } from 'playwright-core/lib/coreBundle';
 import { ManualPromise } from '@isomorphic/manualPromise';
 import { isUnderTest } from '@utils/debug';
+import { isPathInside } from '@utils/fileUtils';
 import { HttpServer } from '@utils/httpServer';
 import { gracefullyProcessExitDoNotHang } from '@utils/processLauncher';
 
@@ -35,7 +37,7 @@ import { wrapReporterAsV2 } from '../reporters/reporterV2';
 import type { Transport } from '@utils/httpServer';
 import type * as reporterTypes from '../../types/testReporter';
 import type { ConfigLocation } from '../common';
-import type { ReportEntry, TestServerInterface, TestServerInterfaceEventEmitters } from '../isomorphic/testServerInterface';
+import type { ReportEntry, TestServerInterface, TestServerInterfaceEventEmitters, UpdateSnapshotParams } from '../isomorphic/testServerInterface';
 import type { ReporterV2 } from '../reporters/reporterV2';
 
 type TraceViewerRedirectOptions = coreServer.TraceViewerRedirectOptions;
@@ -49,6 +51,27 @@ const originalStderrWrite = process.stderr.write;
 
 const originalStdinIsTTY = process.stdin.isTTY;
 
+function allowedFileRoots(configLocation: ConfigLocation, testRunner?: TestRunner): string[] {
+  const roots = new Set<string>([process.cwd(), configLocation.configDir]);
+  const config = testRunner?.lastLoadedConfig();
+  if (config) {
+    for (const project of config.projects) {
+      roots.add(project.project.outputDir);
+      roots.add(project.project.snapshotDir);
+      roots.add(project.project.testDir);
+    }
+  }
+  return [...roots];
+}
+
+function testResultKey(testId: string, resultId: string): string {
+  return JSON.stringify([testId, resultId]);
+}
+
+function attachmentKey(attachment: UpdateSnapshotParams['actual']): string {
+  return JSON.stringify([attachment.name, attachment.contentType]);
+}
+
 class TestServer {
   private _configLocation: ConfigLocation;
   private _configCLIOverrides: ipc.ConfigCLIOverrides;
@@ -70,15 +93,7 @@ class TestServer {
   }
 
   private _allowedFileRoots(): string[] {
-    const roots = new Set<string>([process.cwd(), this._configLocation.configDir]);
-    const config = this._dispatcher?._testRunner.lastLoadedConfig();
-    if (config) {
-      for (const project of config.projects) {
-        roots.add(project.project.outputDir);
-        roots.add(project.project.testDir);
-      }
-    }
-    return [...roots];
+    return allowedFileRoots(this._configLocation, this._dispatcher?._testRunner);
   }
 
   async stop() {
@@ -112,13 +127,17 @@ export type RunTestsParams = {
 
 export class TestServerDispatcher implements TestServerInterface {
   readonly transport: Transport;
+  private _configLocation: ConfigLocation;
   private _serializer: string | undefined;
   private _closeOnDisconnect = false;
+  private _testResultIdForTestId = new Map<string, string>();
+  private _pathForAttachmentForTestResult = new Map</* testResultKey */ string, Map</* attachmentKey */ string, string | null | undefined>>();
   _testRunner: TestRunner;
   private _globalSetupReport: ReportEntry[] | undefined;
   readonly _dispatchEvent: TestServerInterfaceEventEmitters['dispatchEvent'];
 
   constructor(configLocation: ConfigLocation, configCLIOverrides: ipc.ConfigCLIOverrides) {
+    this._configLocation = configLocation;
     this._testRunner = new TestRunner(configLocation, configCLIOverrides);
     this.transport = {
       onconnect: () => {},
@@ -134,8 +153,41 @@ export class TestServerDispatcher implements TestServerInterface {
     this._testRunner.on(TestRunnerEvent.TestPaused, params => this._dispatchEvent('testPaused', { errors: params.errors }));
   }
 
-  private async _wireReporter(messageSink: (message: any) => void) {
-    return await createReporterForTestServer(this._serializer, messageSink);
+  private async _wireReporter(messageSink: (message: ReportEntry) => void) {
+    return await createReporterForTestServer(this._serializer, message => {
+      this._trackAttachmentPaths(message);
+      messageSink(message);
+    });
+  }
+
+  private _trackAttachmentPaths(message: ReportEntry) {
+    switch (message.method) {
+      case 'onTestBegin': {
+        const { testId, result } = message.params;
+        const previousResultId = this._testResultIdForTestId.get(testId);
+        if (previousResultId)
+          this._pathForAttachmentForTestResult.delete(testResultKey(testId, previousResultId));
+        this._testResultIdForTestId.set(testId, result.id);
+        return;
+      }
+
+      case 'onAttach': {
+        if (this._testResultIdForTestId.get(message.params.testId) !== message.params.resultId)
+          throw new Error(`Unknown test result: ${message.params.resultId}`);
+        const resultKey = testResultKey(message.params.testId, message.params.resultId);
+        let paths = this._pathForAttachmentForTestResult.get(resultKey);
+        if (!paths) {
+          paths = new Map();
+          this._pathForAttachmentForTestResult.set(resultKey, paths);
+        }
+        for (const attachment of message.params.attachments) {
+          // Attachments without paths can also make a name ambiguous.
+          const key = attachmentKey(attachment);
+          paths.set(key, paths.has(key) ? null : attachment.path);
+        }
+        return;
+      }
+    }
   }
 
   private async _collectingReporter(): Promise<{ reporter: ReporterV2, report: ReportEntry[] }> {
@@ -195,6 +247,26 @@ export class TestServerDispatcher implements TestServerInterface {
     await this._testRunner.clearCache();
   }
 
+  async updateSnapshot(params: Parameters<TestServerInterface['updateSnapshot']>[0]): ReturnType<TestServerInterface['updateSnapshot']> {
+    const config = this._testRunner.lastLoadedConfig();
+    if (!config)
+      throw new Error('Cannot update snapshot before loading the configuration');
+    const actualName = params.actual.name.match(/^(.*)-actual\.png$/)?.[1];
+    const expectedName = params.expected.name.match(/^(.*)-expected\.png$/)?.[1];
+    if (!actualName || actualName !== expectedName)
+      throw new Error('Actual and expected attachments do not form a snapshot pair');
+    const pathForAttachment = this._pathForAttachmentForTestResult.get(testResultKey(params.testId, params.resultId));
+    const actualPath = pathForAttachment?.get(attachmentKey(params.actual));
+    const expectedPath = pathForAttachment?.get(attachmentKey(params.expected));
+    if (!actualPath || !expectedPath)
+      throw new Error('Snapshot attachments are not registered or have duplicates for this test result');
+    if (!config.projects.some(project => isPathInside(project.project.outputDir, actualPath)))
+      throw new Error('Actual snapshot path is outside of the configured output directories');
+    if (!allowedFileRoots(this._configLocation, this._testRunner).some(root => isPathInside(root, expectedPath)))
+      throw new Error('Expected snapshot path is outside of the allowed file roots');
+    await fs.promises.copyFile(actualPath, expectedPath);
+  }
+
   async listFiles(params: Parameters<TestServerInterface['listFiles']>[0]): ReturnType<TestServerInterface['listFiles']> {
     const { reporter, report } = await this._collectingReporter();
     const { status } = await this._testRunner.listFiles(reporter, params.projects);
@@ -348,7 +420,7 @@ function chunkToPayload(type: 'stdout' | 'stderr', chunk: Buffer | string): Stdi
   return { type, text: chunk };
 }
 
-async function createReporterForTestServer(file: string | undefined, messageSink: (message: any) => void): Promise<ReporterV2> {
+async function createReporterForTestServer(file: string | undefined, messageSink: (message: ReportEntry) => void): Promise<ReporterV2> {
   const reporterConstructor = file ? await loadReporter(null, file) : UIModeReporter;
   return wrapReporterAsV2(new reporterConstructor({
     _send: messageSink,

packages/trace-viewer/src/ui/attachmentsTab.css

@@ -36,6 +36,20 @@
   margin-top: 10px;
 }
 
+.attachments-image-diff-header {
+  position: relative;
+}
+
+.attachments-update-snapshot {
+  background-color: var(--vscode-editor-inactiveSelectionBackground);
+  font-weight: normal;
+  padding: 4px 12px;
+  position: absolute;
+  right: 5px;
+  text-transform: none;
+  top: 5px;
+}
+
 .attachment-item {
   margin: 4px 8px;
 }

packages/trace-viewer/src/ui/attachmentsTab.tsx

@@ -22,11 +22,16 @@ import { CodeMirrorWrapper, lineHeight } from '@web/components/codeMirrorWrapper
 import { isTextualMimeType } from '@isomorphic/mimeType';
 import { Expandable } from '@web/components/expandable';
 import { linkifyText } from '@web/renderUtils';
+import { ToolbarButton } from '@web/components/toolbarButton';
 import { clsx, useFlash } from '@web/uiUtils';
 import { useTraceModel } from './traceModelContext';
 
 import type { Attachment, TraceModel } from '@isomorphic/trace/traceModel';
 
+type SnapshotAttachment = Pick<Attachment, 'contentType' | 'name'>;
+
+export type UpdateSnapshot = (params: { actual: SnapshotAttachment, expected: SnapshotAttachment }) => Promise<void>;
+
 type ExpandableAttachmentProps = {
   attachment: Attachment;
   reveal: any;
@@ -90,9 +95,47 @@ const ExpandableAttachment: React.FunctionComponent<ExpandableAttachmentProps> =
   </div>;
 };
 
+function UpdateSnapshotButton({ actual, expected, onUpdateSnapshot }: {
+  actual: SnapshotAttachment,
+  expected: SnapshotAttachment,
+  onUpdateSnapshot: UpdateSnapshot,
+}) {
+  const [saving, setSaving] = React.useState(false);
+  const [saved, triggerSavedFlash] = useFlash();
+  const [error, setError] = React.useState<string>();
+
+  const updateSnapshot = React.useCallback(async () => {
+    setSaving(true);
+    setError(undefined);
+    try {
+      await onUpdateSnapshot({
+        actual: { name: actual.name, contentType: actual.contentType },
+        expected: { name: expected.name, contentType: expected.contentType },
+      });
+      triggerSavedFlash();
+    } catch (error) {
+      setError(error instanceof Error ? error.message : String(error));
+    } finally {
+      setSaving(false);
+    }
+  }, [actual, expected, onUpdateSnapshot, triggerSavedFlash]);
+
+  const label = error ? 'Retry save' : 'Save actual as expected';
+
+  return <ToolbarButton
+    className='attachments-update-snapshot'
+    disabled={saving}
+    errorBadge={error}
+    icon={saved ? 'check' : undefined}
+    onClick={updateSnapshot}
+    title={error || label}
+  >{label}</ToolbarButton>;
+}
+
 export const AttachmentsTab: React.FunctionComponent<{
   revealedAttachmentCallId?: { callId: string },
-}> = ({ revealedAttachmentCallId }) => {
+  onUpdateSnapshot?: UpdateSnapshot,
+}> = ({ revealedAttachmentCallId, onUpdateSnapshot }) => {
   const model = useTraceModel();
   const { diffMap, screenshots, attachments } = React.useMemo(() => {
     const attachments = new Set(model?.visibleAttachments ?? []);
@@ -122,16 +165,25 @@ export const AttachmentsTab: React.FunctionComponent<{
     return <PlaceholderPanel text='No attachments' />;
 
   return <div className='attachments-tab'>
-    {[...diffMap.values()].map(({ expected, actual, diff }) => {
-      return <>
-        {expected && actual && <div className='attachments-section'>Image diff</div>}
-        {expected && actual && <ImageDiffView noTargetBlank={true} diff={{
+    {[...diffMap.entries()].map(([name, { expected, actual, diff }]) => {
+      if (!expected || !actual)
+        return null;
+      return <React.Fragment key={`${name}-${actual.callId}`}>
+        <div className={clsx('attachments-section', onUpdateSnapshot && 'attachments-image-diff-header')}>
+          <span>Image diff</span>
+          {onUpdateSnapshot && <UpdateSnapshotButton
+            actual={actual}
+            expected={expected}
+            onUpdateSnapshot={onUpdateSnapshot}
+          />}
+        </div>
+        <ImageDiffView noTargetBlank={true} diff={{
           name: 'Image diff',
           expected: { attachment: { ...expected, path: downloadURL(model, expected) }, title: 'Expected' },
           actual: { attachment: { ...actual, path: downloadURL(model, actual) } },
           diff: diff ? { attachment: { ...diff, path: downloadURL(model, diff) } } : undefined,
-        }} />}
-      </>;
+        }} />
+      </React.Fragment>;
     })}
     {screenshots.size ? <div className='attachments-section'>Screenshots</div> : undefined}
     {[...screenshots.values()].map((a, i) => {

packages/trace-viewer/src/ui/uiModeTraceView.tsx

@@ -15,6 +15,7 @@
  */
 
 import { artifactsFolderName } from '@testIsomorphic/folders';
+import { TeleTestResult } from '@testIsomorphic/teleReceiver';
 import type { TreeItem } from '@testIsomorphic/testTree';
 import '@web/common.css';
 import '@web/third_party/vscode/codicon.css';
@@ -23,17 +24,20 @@ import React from 'react';
 import type { ContextEntry } from '@isomorphic/trace/entries';
 import type { SourceLocation } from '@isomorphic/trace/traceModel';
 import { TraceModel } from '@isomorphic/trace/traceModel';
+import type { UpdateSnapshotParams } from '@testIsomorphic/testServerInterface';
+import type { UpdateSnapshot } from './attachmentsTab';
 import { Workbench } from './workbench';
 
 export const TraceView: React.FC<{
   item: { treeItem?: TreeItem, testFile?: SourceLocation, testCase?: reporterTypes.TestCase },
   rootDir?: string,
   onOpenExternally?: (location: SourceLocation) => void,
+  onUpdateSnapshot?: (params: UpdateSnapshotParams) => Promise<void>,
   revealSource?: boolean,
   pathSeparator: string,
   onModelChange?: (model: TraceModel | undefined) => void,
-}> = ({ item, rootDir, onOpenExternally, revealSource, pathSeparator, onModelChange }) => {
-  const [model, setModel] = React.useState<{ model: TraceModel, isLive: boolean } | undefined>(undefined);
+}> = ({ item, rootDir, onOpenExternally, onUpdateSnapshot, revealSource, pathSeparator, onModelChange }) => {
+  const [model, setModel] = React.useState<{ model: TraceModel, isLive: boolean, result?: reporterTypes.TestResult } | undefined>(undefined);
   const [counter, setCounter] = React.useState(0);
   const pollTimer = React.useRef<NodeJS.Timeout | null>(null);
 
@@ -55,7 +59,7 @@ export const TraceView: React.FC<{
     // Test finished.
     const attachment = result && result.duration >= 0 && result.attachments.find(a => a.name === 'trace');
     if (attachment && attachment.path) {
-      loadSingleTraceFile(attachment.path, result.startTime.getTime()).then(model => setModel({ model, isLive: false }));
+      loadSingleTraceFile(attachment.path, result.startTime.getTime()).then(model => setModel({ model, isLive: false, result }));
       return;
     }
 
@@ -93,6 +97,14 @@ export const TraceView: React.FC<{
     onModelChange?.(model?.model);
   }, [model, onModelChange]);
 
+  const testCase = item.testCase;
+  const testResult = testCase?.results[0];
+  const updateSnapshot: UpdateSnapshot | undefined = testCase && testResult instanceof TeleTestResult && model?.result === testResult && onUpdateSnapshot ? params => onUpdateSnapshot({
+    ...params,
+    testId: testCase.id,
+    resultId: testResult._id,
+  }) : undefined;
+
   return <Workbench
     model={model?.model}
     key='workbench'
@@ -103,6 +115,7 @@ export const TraceView: React.FC<{
     status={item.treeItem?.status}
     defaultAnnotations={item.testCase?.annotations ?? []}
     onOpenExternally={onOpenExternally}
+    onUpdateSnapshot={updateSnapshot}
     revealSource={revealSource}
   />;
 };

packages/trace-viewer/src/ui/uiModeView.tsx

@@ -484,6 +484,7 @@ export const UIModeView: React.FC<{}> = ({
             rootDir={testModel?.config?.rootDir}
             revealSource={revealSource}
             onOpenExternally={location => testServerConnection?.openNoReply({ location: { file: location.file, line: location.line, column: location.column } })}
+            onUpdateSnapshot={testServerConnection ? params => testServerConnection.updateSnapshot(params) : undefined}
           />
         </div>
       </div>}