chore: allow matching aria snapshot in trace viewer (#34302)

package-lock.json

@@ -7964,7 +7964,10 @@
       }
     },
     "packages/trace-viewer": {
-      "version": "0.0.0"
+      "version": "0.0.0",
+      "dependencies": {
+        "yaml": "^2.6.0"
+      }
     },
     "packages/web": {
       "version": "0.0.0",

packages/playwright-core/src/server/injected/ariaSnapshot.ts

@@ -241,7 +241,7 @@ function matchesNode(node: AriaNode | string, template: AriaTemplateNode, depth:
   if (typeof node === 'string' && template.kind === 'text')
     return matchesTextNode(node, template);
 
-  if (typeof node === 'object' && template.kind === 'role') {
+  if (node !== null && typeof node === 'object' && template.kind === 'role') {
     if (template.role !== 'fragment' && template.role !== node.role)
       return false;
     if (template.checked !== undefined && template.checked !== node.checked)
@@ -285,20 +285,22 @@ function containsList(children: (AriaNode | string)[], template: AriaTemplateNod
 
 function matchesNodeDeep(root: AriaNode, template: AriaTemplateNode, collectAll: boolean): AriaNode[] {
   const results: AriaNode[] = [];
-  const visit = (node: AriaNode | string): boolean => {
+  const visit = (node: AriaNode | string, parent: AriaNode | null): boolean => {
     if (matchesNode(node, template, 0)) {
-      results.push(node as AriaNode);
+      const result = typeof node === 'string' ? parent : node;
+      if (result)
+        results.push(result);
       return !collectAll;
     }
     if (typeof node === 'string')
       return false;
     for (const child of node.children || []) {
-      if (visit(child))
+      if (visit(child, node))
         return true;
     }
     return false;
   };
-  visit(root);
+  visit(root, null);
   return results;
 }
 

packages/playwright-core/src/utils/isomorphic/ariaSnapshot.ts

@@ -159,6 +159,15 @@ export function parseAriaSnapshot(yaml: YamlLibrary, text: string, options: yaml
       // - role "name": "text"
       const valueIsScalar = value instanceof yaml.Scalar;
       if (valueIsScalar) {
+        const type = typeof value.value;
+        if (type !== 'string' && type !== 'number' && type !== 'boolean') {
+          errors.push({
+            message: 'Node value should be a string or a sequence',
+            range: convertRange(((entry.value as any).range || map.range)),
+          });
+          continue;
+        }
+
         container.children.push({
           ...childNode,
           children: [{
@@ -193,7 +202,7 @@ export function parseAriaSnapshot(yaml: YamlLibrary, text: string, options: yaml
   if (!(yamlDoc.contents instanceof yaml.YAMLSeq)) {
     errors.push({
       message: 'Aria snapshot must be a YAML sequence, elements starting with " -"',
-      range: convertRange(yamlDoc.contents!.range),
+      range: yamlDoc.contents ? convertRange(yamlDoc.contents!.range) : [{ line: 0, col: 0 }, { line: 0, col: 0 }],
     });
   }
   if (errors.length)
@@ -214,7 +223,7 @@ function normalizeWhitespace(text: string) {
 }
 
 export function valueOrRegex(value: string): string | AriaRegex {
-  return value.startsWith('/') && value.endsWith('/') ? { pattern: value.slice(1, -1) } : normalizeWhitespace(value);
+  return value.startsWith('/') && value.endsWith('/') && value.length > 1 ? { pattern: value.slice(1, -1) } : normalizeWhitespace(value);
 }
 
 export class KeyParser {

packages/trace-viewer/package.json

@@ -2,5 +2,8 @@
   "name": "trace-viewer",
   "private": true,
   "version": "0.0.0",
-  "type": "module"
+  "type": "module",
+  "dependencies": {
+    "yaml": "^2.6.0"
+  }
 }

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

@@ -15,30 +15,58 @@
  */
 
 import { CodeMirrorWrapper } from '@web/components/codeMirrorWrapper';
-import type { Language } from '@web/components/codeMirrorWrapper';
+import type { Language, SourceHighlight } from '@web/components/codeMirrorWrapper';
 import { ToolbarButton } from '@web/components/toolbarButton';
 import { copy } from '@web/uiUtils';
 import * as React from 'react';
+import type { HighlightedElement } from './snapshotTab';
 import './sourceTab.css';
+import { parseAriaSnapshot } from '@isomorphic/ariaSnapshot';
+import yaml from 'yaml';
 
 export const InspectorTab: React.FunctionComponent<{
   sdkLanguage: Language,
   setIsInspecting: (isInspecting: boolean) => void,
-  highlightedLocator: string,
+  highlightedElement: HighlightedElement,
-  setHighlightedLocator: (locator: string) => void,
+  setHighlightedElement: (element: HighlightedElement) => void,
-}> = ({ sdkLanguage, setIsInspecting, highlightedLocator, setHighlightedLocator }) => {
+}> = ({ sdkLanguage, setIsInspecting, highlightedElement, setHighlightedElement }) => {
+  const [ariaSnapshotErrors, setAriaSnapshotErrors] = React.useState<SourceHighlight[]>();
+  const onAriaEditorChange = React.useCallback((ariaSnapshot: string) => {
+    const { errors } = parseAriaSnapshot(yaml, ariaSnapshot, { prettyErrors: false });
+    const highlights = errors.map(error => {
+      const highlight: SourceHighlight = {
+        message: error.message,
+        line: error.range[1].line,
+        column: error.range[1].col,
+        type: 'subtle-error',
+      };
+      return highlight;
+    });
+    setAriaSnapshotErrors(highlights);
+    setHighlightedElement({ ...highlightedElement, ariaSnapshot, lastEdited: 'ariaSnapshot' });
+    setIsInspecting(false);
+  }, [highlightedElement, setHighlightedElement, setIsInspecting]);
+
   return <div className='vbox' style={{ backgroundColor: 'var(--vscode-sideBar-background)' }}>
     <div style={{ margin: '10px 0px 10px 10px', color: 'var(--vscode-editorCodeLens-foreground)', flex: 'none' }}>Locator</div>
     <div style={{ margin: '0 10px 10px', flex: 'auto' }}>
-      <CodeMirrorWrapper text={highlightedLocator} language={sdkLanguage} focusOnChange={true} isFocused={true} wrapLines={true} onChange={text => {
+      <CodeMirrorWrapper text={highlightedElement.locator || ''} language={sdkLanguage} isFocused={true} wrapLines={true} onChange={text => {
         // Updating text needs to go first - react can squeeze a render between the state updates.
-        setHighlightedLocator(text);
+        setHighlightedElement({ ...highlightedElement, locator: text, lastEdited: 'locator' });
         setIsInspecting(false);
-      }}></CodeMirrorWrapper>
+      }} />
+    </div>
+    <div style={{ margin: '10px 0px 10px 10px', color: 'var(--vscode-editorCodeLens-foreground)', flex: 'none' }}>Aria</div>
+    <div style={{ margin: '0 10px 10px', flex: 'auto' }}>
+      <CodeMirrorWrapper
+        text={highlightedElement.ariaSnapshot || ''}
+        wrapLines={false}
+        highlight={ariaSnapshotErrors}
+        onChange={onAriaEditorChange} />
     </div>
     <div style={{ position: 'absolute', right: 5, top: 5 }}>
       <ToolbarButton icon='files' title='Copy locator' onClick={() => {
-        copy(highlightedLocator);
+        copy(highlightedElement.locator || '');
       }}></ToolbarButton>
     </div>
   </div>;

packages/trace-viewer/src/ui/recorder/recorderView.tsx

@@ -37,6 +37,7 @@ import { ActionListView } from './actionListView';
 import { BackendContext, BackendProvider } from './backendContext';
 import type { Language } from '@isomorphic/locatorGenerators';
 import { SettingsToolbarButton } from '../settingsToolbarButton';
+import type { HighlightedElement } from '../snapshotTab';
 
 export const RecorderView: React.FunctionComponent = () => {
   const searchParams = new URLSearchParams(window.location.search);
@@ -56,8 +57,8 @@ export const Workbench: React.FunctionComponent = () => {
   const [fileId, setFileId] = React.useState<string | undefined>();
   const [selectedStartTime, setSelectedStartTime] = React.useState<number | undefined>(undefined);
   const [isInspecting, setIsInspecting] = React.useState(false);
-  const [highlightedLocatorInProperties, setHighlightedLocatorInProperties] = React.useState<string>('');
+  const [highlightedElementInProperties, setHighlightedElementInProperties] = React.useState<HighlightedElement>({ lastEdited: 'none' });
-  const [highlightedLocatorInTrace, setHighlightedLocatorInTrace] = React.useState<string>('');
+  const [highlightedElementInTrace, setHighlightedElementInTrace] = React.useState<HighlightedElement>({ lastEdited: 'none' });
   const [traceCallId, setTraceCallId] = React.useState<string | undefined>();
 
   const setSelectedAction = React.useCallback((action: actionTypes.ActionInContext | undefined) => {
@@ -103,15 +104,15 @@ export const Workbench: React.FunctionComponent = () => {
     return { boundaries };
   }, [model]);
 
-  const locatorPickedInTrace = React.useCallback((locator: string) => {
+  const elementPickedInTrace = React.useCallback((element: HighlightedElement) => {
-    setHighlightedLocatorInProperties(locator);
+    setHighlightedElementInProperties(element);
-    setHighlightedLocatorInTrace('');
+    setHighlightedElementInTrace({ lastEdited: 'none' });
     setIsInspecting(false);
   }, []);
 
-  const locatorTypedInProperties = React.useCallback((locator: string) => {
+  const elementTypedInProperties = React.useCallback((element: HighlightedElement) => {
-    setHighlightedLocatorInTrace(locator);
+    setHighlightedElementInTrace(element);
-    setHighlightedLocatorInProperties(locator);
+    setHighlightedElementInProperties(element);
   }, []);
 
   const actionList = <ActionListView
@@ -157,14 +158,14 @@ export const Workbench: React.FunctionComponent = () => {
     callId={traceCallId}
     isInspecting={isInspecting}
     setIsInspecting={setIsInspecting}
-    highlightedLocator={highlightedLocatorInTrace}
+    highlightedElement={highlightedElementInTrace}
-    setHighlightedLocator={locatorPickedInTrace} />;
+    setHighlightedElement={elementPickedInTrace} />;
   const propertiesView = <PropertiesView
     sdkLanguage={sdkLanguage}
     boundaries={boundaries}
     setIsInspecting={setIsInspecting}
-    highlightedLocator={highlightedLocatorInProperties}
+    highlightedElement={highlightedElementInProperties}
-    setHighlightedLocator={locatorTypedInProperties}
+    setHighlightedElement={elementTypedInProperties}
     sourceLocation={sourceLocation} />;
 
   return <div className='vbox workbench'>
@@ -192,15 +193,15 @@ const PropertiesView: React.FunctionComponent<{
   sdkLanguage: Language,
   boundaries: Boundaries,
   setIsInspecting: (value: boolean) => void,
-  highlightedLocator: string,
+  highlightedElement: HighlightedElement,
-  setHighlightedLocator: (locator: string) => void,
+  setHighlightedElement: (element: HighlightedElement) => void,
   sourceLocation: modelUtil.SourceLocation | undefined,
 }> = ({
   sdkLanguage,
   boundaries,
   setIsInspecting,
-  highlightedLocator,
+  highlightedElement,
-  setHighlightedLocator,
+  setHighlightedElement,
   sourceLocation,
 }) => {
   const model = React.useContext(ModelContext);
@@ -215,8 +216,8 @@ const PropertiesView: React.FunctionComponent<{
     render: () => <InspectorTab
       sdkLanguage={sdkLanguage}
       setIsInspecting={setIsInspecting}
-      highlightedLocator={highlightedLocator}
+      highlightedElement={highlightedElement}
-      setHighlightedLocator={setHighlightedLocator} />,
+      setHighlightedElement={setHighlightedElement} />,
   };
 
   const sourceTab: TabbedPaneTabModel = {
@@ -260,15 +261,15 @@ const TraceView: React.FunctionComponent<{
   callId: string | undefined,
   isInspecting: boolean;
   setIsInspecting: (value: boolean) => void;
-  highlightedLocator: string;
+  highlightedElement: HighlightedElement;
-  setHighlightedLocator: (locator: string) => void;
+  setHighlightedElement: (element: HighlightedElement) => void;
 }> = ({
   sdkLanguage,
   callId,
   isInspecting,
   setIsInspecting,
-  highlightedLocator,
+  highlightedElement,
-  setHighlightedLocator,
+  setHighlightedElement,
 }) => {
   const model = React.useContext(ModelContext);
 
@@ -292,7 +293,7 @@ const TraceView: React.FunctionComponent<{
     testIdAttributeName='data-testid'
     isInspecting={isInspecting}
     setIsInspecting={setIsInspecting}
-    highlightedLocator={highlightedLocator}
+    highlightedElement={highlightedElement}
-    setHighlightedLocator={setHighlightedLocator}
+    setHighlightedElement={setHighlightedElement}
     snapshotUrls={snapshotUrls} />;
 };

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

@@ -30,6 +30,14 @@ import { locatorOrSelectorAsSelector } from '@isomorphic/locatorParser';
 import { TabbedPaneTab } from '@web/components/tabbedPane';
 import { BrowserFrame } from './browserFrame';
 import type { ElementInfo } from '@recorder/recorderTypes';
+import { parseAriaSnapshot } from '@isomorphic/ariaSnapshot';
+import yaml from 'yaml';
+
+export type HighlightedElement = {
+  locator?: string,
+  ariaSnapshot?: string
+  lastEdited: 'locator' | 'ariaSnapshot' | 'none';
+};
 
 export const SnapshotTabsView: React.FunctionComponent<{
   action: ActionTraceEvent | undefined,
@@ -38,9 +46,9 @@ export const SnapshotTabsView: React.FunctionComponent<{
   testIdAttributeName: string,
   isInspecting: boolean,
   setIsInspecting: (isInspecting: boolean) => void,
-  highlightedLocator: string,
+  highlightedElement: HighlightedElement,
-  setHighlightedLocator: (locator: string) => void,
+  setHighlightedElement: (element: HighlightedElement) => void,
-}> = ({ action, sdkLanguage, testIdAttributeName, isInspecting, setIsInspecting, highlightedLocator, setHighlightedLocator }) => {
+}> = ({ action, sdkLanguage, testIdAttributeName, isInspecting, setIsInspecting, highlightedElement, setHighlightedElement }) => {
   const [snapshotTab, setSnapshotTab] = React.useState<'action'|'before'|'after'>('action');
 
   // eslint-disable-next-line @typescript-eslint/no-unused-vars
@@ -81,8 +89,8 @@ export const SnapshotTabsView: React.FunctionComponent<{
       testIdAttributeName={testIdAttributeName}
       isInspecting={isInspecting}
       setIsInspecting={setIsInspecting}
-      highlightedLocator={highlightedLocator}
+      highlightedElement={highlightedElement}
-      setHighlightedLocator={setHighlightedLocator}
+      setHighlightedElement={setHighlightedElement}
     />
   </div>;
 };
@@ -93,9 +101,9 @@ export const SnapshotView: React.FunctionComponent<{
   testIdAttributeName: string,
   isInspecting: boolean,
   setIsInspecting: (isInspecting: boolean) => void,
-  highlightedLocator: string,
+  highlightedElement: HighlightedElement,
-  setHighlightedLocator: (locator: string) => void,
+  setHighlightedElement: (element: HighlightedElement) => void,
-}> = ({ snapshotUrls, sdkLanguage, testIdAttributeName, isInspecting, setIsInspecting, highlightedLocator, setHighlightedLocator }) => {
+}> = ({ snapshotUrls, sdkLanguage, testIdAttributeName, isInspecting, setIsInspecting, highlightedElement, setHighlightedElement }) => {
   const iframeRef0 = React.useRef<HTMLIFrameElement>(null);
   const iframeRef1 = React.useRef<HTMLIFrameElement>(null);
   const [snapshotInfo, setSnapshotInfo] = React.useState<SnapshotInfo>({ viewport: kDefaultViewport, url: '' });
@@ -158,16 +166,16 @@ export const SnapshotView: React.FunctionComponent<{
       isInspecting={isInspecting}
       sdkLanguage={sdkLanguage}
       testIdAttributeName={testIdAttributeName}
-      highlightedLocator={highlightedLocator}
+      highlightedElement={highlightedElement}
-      setHighlightedLocator={setHighlightedLocator}
+      setHighlightedElement={setHighlightedElement}
       iframe={iframeRef0.current}
       iteration={loadingRef.current.iteration} />
     <InspectModeController
       isInspecting={isInspecting}
       sdkLanguage={sdkLanguage}
       testIdAttributeName={testIdAttributeName}
-      highlightedLocator={highlightedLocator}
+      highlightedElement={highlightedElement}
-      setHighlightedLocator={setHighlightedLocator}
+      setHighlightedElement={setHighlightedElement}
       iframe={iframeRef1.current}
       iteration={loadingRef.current.iteration} />
     <SnapshotWrapper snapshotInfo={snapshotInfo}>
@@ -223,10 +231,10 @@ export const InspectModeController: React.FunctionComponent<{
   isInspecting: boolean,
   sdkLanguage: Language,
   testIdAttributeName: string,
-  highlightedLocator: string,
+  highlightedElement: HighlightedElement,
-  setHighlightedLocator: (locator: string) => void,
+  setHighlightedElement: (element: HighlightedElement) => void,
   iteration: number,
-}> = ({ iframe, isInspecting, sdkLanguage, testIdAttributeName, highlightedLocator, setHighlightedLocator, iteration }) => {
+}> = ({ iframe, isInspecting, sdkLanguage, testIdAttributeName, highlightedElement, setHighlightedElement, iteration }) => {
   React.useEffect(() => {
     const recorders: { recorder: Recorder, frameSelector: string }[] = [];
     const isUnderTest = new URLSearchParams(window.location.search).get('isUnderTest') === 'true';
@@ -236,17 +244,25 @@ export const InspectModeController: React.FunctionComponent<{
       // Potential cross-origin exceptions.
     }
 
+    const parsedSnapshot = highlightedElement.lastEdited === 'ariaSnapshot' && highlightedElement.ariaSnapshot ? parseAriaSnapshot(yaml, highlightedElement.ariaSnapshot) : undefined;
+    const fullSelector = highlightedElement.lastEdited === 'locator' && highlightedElement.locator ? locatorOrSelectorAsSelector(sdkLanguage, highlightedElement.locator, testIdAttributeName) : undefined;
     for (const { recorder, frameSelector } of recorders) {
-      const actionSelector = locatorOrSelectorAsSelector(sdkLanguage, highlightedLocator, testIdAttributeName);
+      const actionSelector = fullSelector?.startsWith(frameSelector) ? fullSelector.substring(frameSelector.length).trim() : undefined;
+      const ariaTemplate = parsedSnapshot?.errors.length === 0 ? parsedSnapshot.fragment : undefined;
       recorder.setUIState({
         mode: isInspecting ? 'inspecting' : 'none',
-        actionSelector: actionSelector.startsWith(frameSelector) ? actionSelector.substring(frameSelector.length).trim() : undefined,
+        actionSelector,
+        ariaTemplate,
         language: sdkLanguage,
         testIdAttributeName,
         overlay: { offsetX: 0 },
       }, {
         async elementPicked(elementInfo: ElementInfo) {
-          setHighlightedLocator(asLocator(sdkLanguage, frameSelector + elementInfo.selector));
+          setHighlightedElement({
+            locator: asLocator(sdkLanguage, frameSelector + elementInfo.selector),
+            ariaSnapshot: elementInfo.ariaSnapshot,
+            lastEdited: 'none',
+          });
         },
         highlightUpdated() {
           for (const r of recorders) {
@@ -256,7 +272,7 @@ export const InspectModeController: React.FunctionComponent<{
         }
       });
     }
-  }, [iframe, isInspecting, highlightedLocator, setHighlightedLocator, sdkLanguage, testIdAttributeName, iteration]);
+  }, [iframe, isInspecting, highlightedElement, setHighlightedElement, sdkLanguage, testIdAttributeName, iteration]);
   return <></>;
 };
 

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

@@ -42,6 +42,7 @@ import './workbench.css';
 import { testStatusIcon, testStatusText } from './testUtils';
 import type { UITestStatus } from './testUtils';
 import type { AfterActionTraceEventAttachment } from '@trace/trace';
+import type { HighlightedElement } from './snapshotTab';
 
 export const Workbench: React.FunctionComponent<{
   model?: modelUtil.MultiTraceModel,
@@ -65,7 +66,7 @@ export const Workbench: React.FunctionComponent<{
   const [selectedNavigatorTab, setSelectedNavigatorTab] = React.useState<string>('actions');
   const [selectedPropertiesTab, setSelectedPropertiesTab] = useSetting<string>('propertiesTab', showSourcesFirst ? 'source' : 'call');
   const [isInspecting, setIsInspectingState] = React.useState(false);
-  const [highlightedLocator, setHighlightedLocator] = React.useState<string>('');
+  const [highlightedElement, setHighlightedElement] = React.useState<HighlightedElement>({ lastEdited: 'none' });
   const [selectedTime, setSelectedTime] = React.useState<Boundaries | undefined>();
   const [sidebarLocation, setSidebarLocation] = useSetting<'bottom' | 'right'>('propertiesSidebarLocation', 'bottom');
 
@@ -140,8 +141,8 @@ export const Workbench: React.FunctionComponent<{
     setIsInspectingState(value);
   }, [setIsInspectingState, selectPropertiesTab, isInspecting]);
 
-  const locatorPicked = React.useCallback((locator: string) => {
+  const elementPicked = React.useCallback((element: HighlightedElement) => {
-    setHighlightedLocator(locator);
+    setHighlightedElement(element);
     selectPropertiesTab('inspector');
   }, [selectPropertiesTab]);
 
@@ -170,8 +171,8 @@ export const Workbench: React.FunctionComponent<{
     render: () => <InspectorTab
       sdkLanguage={sdkLanguage}
       setIsInspecting={setIsInspecting}
-      highlightedLocator={highlightedLocator}
+      highlightedElement={highlightedElement}
-      setHighlightedLocator={setHighlightedLocator} />,
+      setHighlightedElement={setHighlightedElement} />,
   };
   const callTab: TabbedPaneTabModel = {
     id: 'call',
@@ -342,8 +343,8 @@ export const Workbench: React.FunctionComponent<{
           testIdAttributeName={model?.testIdAttributeName || 'data-testid'}
           isInspecting={isInspecting}
           setIsInspecting={setIsInspecting}
-          highlightedLocator={highlightedLocator}
+          highlightedElement={highlightedElement}
-          setHighlightedLocator={locatorPicked} />}
+          setHighlightedElement={elementPicked} />}
         sidebar={
           <TabbedPane
             tabs={[actionsTab, metadataTab]}

tests/library/trace-viewer.spec.ts

@@ -24,7 +24,7 @@ import path from 'path';
 import { pathToFileURL } from 'url';
 import { expect, playwrightTest } from '../config/browserTest';
 import type { FrameLocator } from '@playwright/test';
-import { rafraf } from 'tests/page/pageTest';
+import { rafraf, roundBox } from 'tests/page/pageTest';
 
 const test = playwrightTest.extend<TraceViewerFixtures>(traceViewerFixtures);
 
@@ -1096,19 +1096,41 @@ test('should pick locator', async ({ page, runAndTrace, server }) => {
   const snapshot = await traceViewer.snapshotFrame('page.setContent');
   await traceViewer.page.getByTitle('Pick locator').click();
   await snapshot.locator('button').click();
-  await expect(traceViewer.page.locator('.cm-wrapper')).toContainText(`getByRole('button', { name: 'Submit' })`);
+  await expect(traceViewer.page.locator('.cm-wrapper').first()).toContainText(`getByRole('button', { name: 'Submit' })`);
+  await expect(traceViewer.page.locator('.cm-wrapper').last()).toContainText(`- button "Submit"`);
 });
 
-test('should update highlight when typing', async ({ page, runAndTrace, server }) => {
+test('should update highlight when typing locator', async ({ page, runAndTrace, server }) => {
   const traceViewer = await runAndTrace(async () => {
     await page.goto(server.EMPTY_PAGE);
     await page.setContent('<button>Submit</button>');
   });
   const snapshot = await traceViewer.snapshotFrame('page.setContent');
   await traceViewer.page.getByText('Locator').click();
-  await traceViewer.page.locator('.CodeMirror').click();
+  await traceViewer.page.locator('.CodeMirror').first().click();
   await traceViewer.page.keyboard.type('button');
-  await expect(snapshot.locator('x-pw-glass')).toBeVisible();
+
+  const buttonBox = roundBox(await snapshot.locator('button').boundingBox());
+  await expect(snapshot.locator('x-pw-highlight')).toBeVisible();
+  await expect.poll(async () => {
+    return roundBox(await snapshot.locator('x-pw-highlight').boundingBox());
+  }).toEqual(buttonBox);
+});
+
+test('should update highlight when typing snapshot', async ({ page, runAndTrace, server }) => {
+  const traceViewer = await runAndTrace(async () => {
+    await page.goto(server.EMPTY_PAGE);
+    await page.setContent('<button>Submit</button>');
+  });
+  const snapshot = await traceViewer.snapshotFrame('page.setContent');
+  await traceViewer.page.getByText('Locator').click();
+  await traceViewer.page.locator('.CodeMirror').last().click();
+  await traceViewer.page.keyboard.type('- button');
+  const buttonBox = roundBox(await snapshot.locator('button').boundingBox());
+  await expect(snapshot.locator('x-pw-highlight')).toBeVisible();
+  await expect.poll(async () => {
+    return roundBox(await snapshot.locator('x-pw-highlight').boundingBox());
+  }).toEqual(buttonBox);
 });
 
 test('should open trace-1.31', async ({ showTraceViewer }) => {
@@ -1239,7 +1261,7 @@ test('should pick locator in iframe', async ({ page, runAndTrace, server }) => {
     await page.evaluate('2+2');
   });
   await traceViewer.page.getByTitle('Pick locator').click();
-  const cmWrapper = traceViewer.page.locator('.cm-wrapper');
+  const cmWrapper = traceViewer.page.locator('.cm-wrapper').first();
 
   const snapshot = await traceViewer.snapshotFrame('page.evaluate');
 
@@ -1279,7 +1301,7 @@ test('should highlight locator in iframe while typing', async ({ page, runAndTra
 
   const snapshot = await traceViewer.snapshotFrame('page.evaluate');
   await traceViewer.page.getByText('Locator').click();
-  await traceViewer.page.locator('.CodeMirror').click();
+  await traceViewer.page.locator('.CodeMirror').first().click();
 
   const locators = [{
     text: `locator('#frame1').contentFrame().getByText('Hello1')`,