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

@@ -57,6 +57,11 @@ export interface TestServerInterface {
 
   clearCache(params: {}): Promise<void>;
 
+  updateSnapshot(params: {
+    actualPath: string;
+    expectedPath: string;
+  }): 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';
 
@@ -49,6 +51,19 @@ 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];
+}
+
 class TestServer {
   private _configLocation: ConfigLocation;
   private _configCLIOverrides: ipc.ConfigCLIOverrides;
@@ -70,15 +85,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 +119,15 @@ export type RunTestsParams = {
 
 export class TestServerDispatcher implements TestServerInterface {
   readonly transport: Transport;
+  private _configLocation: ConfigLocation;
   private _serializer: string | undefined;
   private _closeOnDisconnect = false;
   _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: () => {},
@@ -195,6 +204,17 @@ 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');
+    if (!config.projects.some(project => isPathInside(project.project.outputDir, params.actualPath)))
+      throw new Error('Actual snapshot path is outside of the configured output directories');
+    if (!allowedFileRoots(this._configLocation, this._testRunner).some(root => isPathInside(root, params.expectedPath)))
+      throw new Error('Expected snapshot path is outside of the allowed file roots');
+    await fs.promises.copyFile(params.actualPath, params.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);

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,14 @@ 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';
 
+export type UpdateSnapshot = (params: { actualPath: string, expectedPath: string }) => Promise<void>;
+
 type ExpandableAttachmentProps = {
   attachment: Attachment;
   reveal: any;
@@ -90,9 +93,45 @@ const ExpandableAttachment: React.FunctionComponent<ExpandableAttachmentProps> =
   </div>;
 };
 
+function UpdateSnapshotButton({ actualPath, expectedPath, onUpdateSnapshot }: {
+  actualPath: string,
+  expectedPath: string,
+  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({ actualPath, expectedPath });
+      triggerSavedFlash();
+    } catch (error) {
+      setError(error instanceof Error ? error.message : String(error));
+    } finally {
+      setSaving(false);
+    }
+  }, [actualPath, expectedPath, 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<{
+  localAttachmentPaths?: Map<Attachment, string>,
   revealedAttachmentCallId?: { callId: string },
-}> = ({ revealedAttachmentCallId }) => {
+  onUpdateSnapshot?: UpdateSnapshot,
+}> = ({ localAttachmentPaths, revealedAttachmentCallId, onUpdateSnapshot }) => {
   const model = useTraceModel();
   const { diffMap, screenshots, attachments } = React.useMemo(() => {
     const attachments = new Set(model?.visibleAttachments ?? []);
@@ -122,16 +161,27 @@ 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;
+      const expectedPath = localAttachmentPaths?.get(expected) || expected.path;
+      const actualPath = localAttachmentPaths?.get(actual) || actual.path;
+      return <React.Fragment key={`${name}-${actual.callId}`}>
+        <div className={clsx('attachments-section', onUpdateSnapshot && 'attachments-image-diff-header')}>
+          <span>Image diff</span>
+          {onUpdateSnapshot && expectedPath && actualPath && <UpdateSnapshotButton
+            actualPath={actualPath}
+            expectedPath={expectedPath}
+            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

@@ -21,19 +21,25 @@ import '@web/third_party/vscode/codicon.css';
 import type * as reporterTypes from 'playwright/types/testReporter';
 import React from 'react';
 import type { ContextEntry } from '@isomorphic/trace/entries';
-import type { SourceLocation } from '@isomorphic/trace/traceModel';
+import type { Attachment, SourceLocation } from '@isomorphic/trace/traceModel';
 import { TraceModel } from '@isomorphic/trace/traceModel';
+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?: UpdateSnapshot,
   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,
+    localAttachmentPaths?: Map<Attachment, string>,
+  } | undefined>(undefined);
   const [counter, setCounter] = React.useState(0);
   const pollTimer = React.useRef<NodeJS.Timeout | null>(null);
 
@@ -55,7 +61,13 @@ 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,
+          localAttachmentPaths: localAttachmentPathsFromResult(model, result),
+        });
+      });
       return;
     }
 
@@ -95,6 +107,7 @@ export const TraceView: React.FC<{
 
   return <Workbench
     model={model?.model}
+    localAttachmentPaths={model?.localAttachmentPaths}
     key='workbench'
     showSourcesFirst={true}
     rootDir={rootDir}
@@ -103,10 +116,32 @@ export const TraceView: React.FC<{
     status={item.treeItem?.status}
     defaultAnnotations={item.testCase?.annotations ?? []}
     onOpenExternally={onOpenExternally}
+    onUpdateSnapshot={onUpdateSnapshot}
     revealSource={revealSource}
   />;
 };
 
+function localAttachmentPathsFromResult(model: TraceModel, testResult: reporterTypes.TestResult): Map<Attachment, string> {
+  const pathsByAttachment = new Map<string, string[]>();
+  for (const attachment of testResult.attachments) {
+    if (!attachment.path)
+      continue;
+    const key = JSON.stringify([attachment.name, attachment.contentType]);
+    const paths = pathsByAttachment.get(key) || [];
+    paths.push(attachment.path);
+    pathsByAttachment.set(key, paths);
+  }
+
+  const localAttachmentPaths = new Map<Attachment, string>();
+  for (const attachment of model.attachments) {
+    const key = JSON.stringify([attachment.name, attachment.contentType]);
+    const path = pathsByAttachment.get(key)?.shift();
+    if (path)
+      localAttachmentPaths.set(attachment, path);
+  }
+  return localAttachmentPaths;
+}
+
 const outputDirForTestCase = (testCase: reporterTypes.TestCase): string | undefined => {
   for (let suite: reporterTypes.Suite | undefined = testCase.parent; suite; suite = suite.parent) {
     if (suite.project())

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>}

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

@@ -21,7 +21,7 @@ import { CallTab } from './callTab';
 import { LogTab } from './logTab';
 import { ErrorsTab, useErrorsTabModel } from './errorsTab';
 import { ConsoleTab, useConsoleTabModel } from './consoleTab';
-import type { TraceModel, SourceLocation, SourceModel } from '@isomorphic/trace/traceModel';
+import type { Attachment, TraceModel, SourceLocation, SourceModel } from '@isomorphic/trace/traceModel';
 import type { ActionEntry } from '@isomorphic/trace/entries';
 import { NetworkTab, useNetworkTabModel } from './networkTab';
 import { SnapshotTabsView } from './snapshotTab';
@@ -32,6 +32,7 @@ import { Timeline } from './timeline';
 import { usePlayback, PlaybackScrubber } from './playbackControl';
 import { MetadataView } from './metadataView';
 import { AttachmentsTab } from './attachmentsTab';
+import type { UpdateSnapshot } from './attachmentsTab';
 import { AnnotationsTab } from './annotationsTab';
 import type { Boundaries } from './geometry';
 import { InspectorTab } from './inspectorTab';
@@ -52,6 +53,7 @@ import type { TreeState } from '@web/components/treeView';
 
 export type WorkbenchProps = {
   model: TraceModel | undefined;
+  localAttachmentPaths?: Map<Attachment, string>;
   showSourcesFirst?: boolean;
   rootDir?: string;
   fallbackLocation?: SourceLocation;
@@ -61,6 +63,7 @@ export type WorkbenchProps = {
   defaultAnnotations?: TestAnnotation[];
   inert?: boolean;
   onOpenExternally?: (location: SourceLocation) => void;
+  onUpdateSnapshot?: UpdateSnapshot;
   revealSource?: boolean;
   testRunMetadata?: MetadataWithCommitInfo;
 };
@@ -73,7 +76,7 @@ export const Workbench: React.FunctionComponent<WorkbenchProps> = props => {
 };
 
 const PartitionedWorkbench: React.FunctionComponent<WorkbenchProps & { partition: string }> = props => {
-  const { partition, model, showSourcesFirst, rootDir, fallbackLocation, isLive, hideTimeline, status, inert, onOpenExternally, revealSource, testRunMetadata } = props;
+  const { partition, model, localAttachmentPaths, showSourcesFirst, rootDir, fallbackLocation, isLive, hideTimeline, status, inert, onOpenExternally, onUpdateSnapshot, revealSource, testRunMetadata } = props;
   // Default annotations come from the test model before the test runs, shown for the empty workbench / trace.
   const annotations = model?.annotations ?? props.defaultAnnotations;
 
@@ -282,7 +285,11 @@ const PartitionedWorkbench: React.FunctionComponent<WorkbenchProps & { partition
     id: 'attachments',
     title: 'Attachments',
     count: model?.visibleAttachments.length,
-    render: () => <AttachmentsTab revealedAttachmentCallId={revealedAttachmentCallId} />
+    render: () => <AttachmentsTab
+      localAttachmentPaths={localAttachmentPaths}
+      revealedAttachmentCallId={revealedAttachmentCallId}
+      onUpdateSnapshot={status === 'failed' ? onUpdateSnapshot : undefined}
+    />
   };
 
   const tabs: TabbedPaneTabModel[] = [