chore: show snapshot for test.step (#35445)

We don't take before/after snapshot for `test.step`. To approximate the snapshots we could take either snapshots from the nested actions or from the outer ones. The current logic is the following:

**beforeSnapshot:**
- `beforeSnapshot` is always taken from the last finished action before the step. It also works nice for the actions without nested actions, such as simple `expect(1).toBe(1);`

**afterSnapshot:**
- We always use `afterSnapshot` from a "nested" action, if there is one. It is exactly what we want for `test.step` and it is acceptable for other actions.
- If there are no "nested" actions, use the `beforeSnapshot` 
  -  works best for simple `expect(a).toBe(b);` case
  - `test.step` without children with snapshot is likely a step with a bunch of `expect(a).toBe(b);` and the same logic as for single expect applies.

Fixes https://github.com/microsoft/playwright/issues/35285

packages/trace-viewer/src/ui/modelUtil.ts

@@ -24,8 +24,9 @@ import type { ActionEntry, ContextEntry, PageEntry } from '../types/entries';
 import type { StackFrame } from '@protocol/channels';
 
 const contextSymbol = Symbol('context');
-const nextInContextSymbol = Symbol('next');
-const prevInListSymbol = Symbol('prev');
+const nextInContextSymbol = Symbol('nextInContext');
+const prevByEndTimeSymbol = Symbol('prevByEndTime');
+const nextByStartTimeSymbol = Symbol('nextByStartTime');
 const eventsSymbol = Symbol('events');
 
 export type SourceLocation = {
@@ -190,6 +191,18 @@ function mergeActionsAndUpdateTiming(contexts: ContextEntry[]) {
     const actions = mergeActionsAndUpdateTimingSameTrace(contexts);
     result.push(...actions);
   }
+
+  result.sort((a1, a2) => {
+    if (a2.parentId === a1.callId)
+      return 1;
+    if (a1.parentId === a2.callId)
+      return -1;
+    return a1.endTime - a2.endTime;
+  });
+
+  for (let i = 1; i < result.length; ++i)
+    (result[i] as any)[prevByEndTimeSymbol] = result[i - 1];
+
   result.sort((a1, a2) => {
     if (a2.parentId === a1.callId)
       return -1;
@@ -198,8 +211,8 @@ function mergeActionsAndUpdateTiming(contexts: ContextEntry[]) {
     return a1.startTime - a2.startTime;
   });
 
-  for (let i = 1; i < result.length; ++i)
-    (result[i] as any)[prevInListSymbol] = result[i - 1];
+  for (let i = 0; i + 1 < result.length; ++i)
+    (result[i] as any)[nextByStartTimeSymbol] = result[i + 1];
 
   return result;
 }
@@ -355,8 +368,12 @@ function nextInContext(action: ActionTraceEvent): ActionTraceEvent {
   return (action as any)[nextInContextSymbol];
 }
 
-export function prevInList(action: ActionTraceEvent): ActionTraceEvent {
-  return (action as any)[prevInListSymbol];
+export function previousActionByEndTime(action: ActionTraceEvent): ActionTraceEvent {
+  return (action as any)[prevByEndTimeSymbol];
+}
+
+export function nextActionByStartTime(action: ActionTraceEvent): ActionTraceEvent {
+  return (action as any)[nextByStartTimeSymbol];
 }
 
 export function stats(action: ActionTraceEvent): { errors: number, warnings: number } {

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

@@ -17,7 +17,7 @@
 import './snapshotTab.css';
 import * as React from 'react';
 import type { ActionTraceEvent } from '@trace/trace';
-import { context, type MultiTraceModel, prevInList } from './modelUtil';
+import { context, type MultiTraceModel, nextActionByStartTime, previousActionByEndTime } from './modelUtil';
 import { Toolbar } from '@web/components/toolbar';
 import { ToolbarButton } from '@web/components/toolbarButton';
 import { clsx, useMeasure, useSetting } from '@web/uiUtils';
@@ -329,14 +329,40 @@ export function collectSnapshots(action: ActionTraceEvent | undefined): Snapshot
   if (!action)
     return {};
 
-  // if the action has no beforeSnapshot, use the last available afterSnapshot.
   let beforeSnapshot: Snapshot | undefined = action.beforeSnapshot ? { action, snapshotName: action.beforeSnapshot } : undefined;
-  let a = action;
-  while (!beforeSnapshot && a) {
-    a = prevInList(a);
-    beforeSnapshot = a?.afterSnapshot ? { action: a, snapshotName: a?.afterSnapshot } : undefined;
+  if (!beforeSnapshot) {
+    // If the action has no beforeSnapshot, use the last available afterSnapshot.
+    for (let a = previousActionByEndTime(action); a; a = previousActionByEndTime(a)) {
+      if (a.endTime <= action.startTime && a.afterSnapshot) {
+        beforeSnapshot = { action: a, snapshotName: a.afterSnapshot };
+        break;
+      }
+    }
   }
-  const afterSnapshot: Snapshot | undefined = action.afterSnapshot ? { action, snapshotName: action.afterSnapshot } : beforeSnapshot;
+
+  let afterSnapshot: Snapshot | undefined = action.afterSnapshot ? { action, snapshotName: action.afterSnapshot } : undefined;
+  if (!afterSnapshot) {
+    let last: ActionTraceEvent | undefined;
+    // - For test.step, we want to use the snapshot of the last nested action.
+    // - For a regular action, we use snapshot of any overlapping in time action
+    //   as a best effort.
+    // - If there are no "nested" actions, use the beforeSnapshot which works best
+    //   for simple `expect(a).toBe(b);` case. Also if the action doesn't have
+    //   afterSnapshot, it likely doesn't have its own beforeSnapshot either,
+    //   and we calculated it above from a previous action.
+    for (let a = nextActionByStartTime(action); a && a.startTime <= action.endTime; a = nextActionByStartTime(a)) {
+      if (a.endTime > action.endTime || !a.afterSnapshot)
+        continue;
+      if (last && last.endTime > a.endTime)
+        continue;
+      last = a;
+    }
+    if (last)
+      afterSnapshot = { action: last, snapshotName: last.afterSnapshot! };
+    else
+      afterSnapshot = beforeSnapshot;
+  }
+
   const actionSnapshot: Snapshot | undefined = action.inputSnapshot ? { action, snapshotName: action.inputSnapshot, hasInputTarget: true } : afterSnapshot;
   if (actionSnapshot)
     actionSnapshot.point = action.point;

tests/playwright-test/ui-mode-trace.spec.ts

@@ -150,6 +150,61 @@ test('should show snapshots for sync assertions', async ({ runUITest }) => {
   ).toHaveText('Submit');
 });
 
+test('should show snapshots for steps', {
+  annotation: { type: 'issue', description: 'https://github.com/microsoft/playwright/issues/35285' }
+}, async ({ runUITest }) => {
+  const { page } = await runUITest({
+    'a.test.ts': `
+      import { test, expect } from '@playwright/test';
+      test.beforeEach(async ({ page }) => {
+        await page.setContent('<div>initial</div>');
+      });
+      test('steps test', async ({ page }) => {
+        await test.step('first', async () => {
+          await page.setContent("<div>foo</div>");
+        });
+        await test.step('middle', async () => {
+          await page.setContent("<div>bar</div>");
+        });
+        await test.step('last', async () => {
+          await page.setContent("<div>baz</div>");
+        });
+      });
+    `,
+  });
+
+  await page.getByText('steps test').dblclick();
+
+  await expect(page.getByTestId('actions-tree')).toMatchAriaSnapshot(`
+    - tree:
+      - treeitem /Before Hooks \\d+[hmsp]+/
+      - treeitem /first \\d+[hmsp]+/
+      - treeitem /middle \\d+[hmsp]+/
+      - treeitem /last \\d+[hmsp]+/
+      - treeitem /After Hooks \\d+[hmsp]+/
+  `);
+
+  await page.getByTestId('actions-tree').getByText('first').click();
+  const snapshot = page.frameLocator('iframe.snapshot-visible[name=snapshot]').locator('div');
+
+  await page.getByText('After', { exact: true }).click();
+  await expect(snapshot).toHaveText('foo');
+  await page.getByText('Before', { exact: true }).click();
+  await expect(snapshot).toHaveText('initial');
+
+  await page.getByTestId('actions-tree').getByText('middle').click();
+  await page.getByText('After', { exact: true }).click();
+  await expect(snapshot).toHaveText('bar');
+  await page.getByText('Before', { exact: true }).click();
+  await expect(snapshot).toHaveText('foo');
+
+  await page.getByTestId('actions-tree').getByText('last').click();
+  await page.getByText('After', { exact: true }).click();
+  await expect(snapshot).toHaveText('baz');
+  await page.getByText('Before', { exact: true }).click();
+  await expect(snapshot).toHaveText('bar');
+});
+
 test('should show image diff', async ({ runUITest }) => {
   const { page } = await runUITest({
     'playwright.config.js': `