fix(steps): make expect.toPass and expect.poll step containers (#30389)

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

packages/playwright-core/src/client/channelOwner.ts

@@ -187,9 +187,7 @@ export abstract class ChannelOwner<T extends channels.Channel = channels.Channel
     try {
       logApiCall(logger, `=> ${apiName} started`, isInternal);
       const apiZone: ApiZone = { apiName, frames, isInternal, reported: false, csi, callCookie, wallTime };
-      const result = await zones.run<ApiZone, Promise<R>>('apiZone', apiZone, async () => {
+      const result = await zones.run('apiZone', apiZone, async () => await func(apiZone));
-        return await func(apiZone);
-      });
       csi?.onApiCallEnd(callCookie);
       logApiCall(logger, `<= ${apiName} succeeded`, isInternal);
       return result;

packages/playwright-core/src/utils/zones.ts

@@ -21,18 +21,18 @@ export type ZoneType = 'apiZone' | 'expectZone' | 'stepZone';
 class ZoneManager {
   private readonly _asyncLocalStorage = new AsyncLocalStorage<Zone<unknown>|undefined>();
 
-  run<T, R>(type: ZoneType, data: T, func: (data: T) => R): R {
+  run<T, R>(type: ZoneType, data: T, func: () => R): R {
     const previous = this._asyncLocalStorage.getStore();
     const zone = new Zone(previous, type, data);
-    return this._asyncLocalStorage.run(zone, () => func(data));
+    return this._asyncLocalStorage.run(zone, func);
   }
 
-  zoneData<T>(type: ZoneType): T | null {
+  zoneData<T>(type: ZoneType): T | undefined {
     for (let zone = this._asyncLocalStorage.getStore(); zone; zone = zone.previous) {
       if (zone.type === type)
         return zone.data as T;
     }
-    return null;
+    return undefined;
   }
 
   exitZones<R>(func: () => R): R {

packages/playwright/src/matchers/expect.ts

@@ -284,9 +284,12 @@ class ExpectMetaInfoProxyHandler implements ProxyHandler<any> {
       };
 
       try {
-        const expectZone: ExpectZone | null = matcherName !== 'toPass' ? { title, wallTime } : null;
         const callback = () => matcher.call(target, ...args);
-        const result = expectZone ? zones.run<ExpectZone, any>('expectZone', expectZone, callback) : callback();
+        // toPass and poll matchers can contain other steps, expects and API calls,
+        // so they behave like a retriable step.
+        const result = (matcherName === 'toPass' || this._info.isPoll) ?
+          zones.run('stepZone', step, callback) :
+          zones.run<ExpectZone, any>('expectZone', { title, wallTime }, callback);
         if (result instanceof Promise)
           return result.then(finalizer).catch(reportStepError);
         finalizer();

packages/playwright/src/worker/testInfo.ts

@@ -252,7 +252,7 @@ export class TestInfoImpl implements TestInfo {
       // Predefined stages form a fixed hierarchy - use the current one as parent.
       parentStep = this._findLastStageStep();
     } else {
-      parentStep = zones.zoneData<TestStepInternal>('stepZone') || undefined;
+      parentStep = zones.zoneData<TestStepInternal>('stepZone');
       if (!parentStep && data.category !== 'test.step') {
         // API steps (but not test.step calls) can be nested by time, instead of by stack.
         // However, do not nest chains of route.continue by checking the title.

tests/playwright-test/test-step.spec.ts

@@ -846,7 +846,6 @@ fixture   |  fixture: context
 
 test('step inside expect.toPass', async ({ runInlineTest }) => {
   test.info().annotations.push({ type: 'issue', description: 'https://github.com/microsoft/playwright/issues/30322' });
-  test.fixme();
   const result = await runInlineTest({
     'reporter.ts': stepIndentReporter,
     'playwright.config.ts': `
@@ -883,7 +882,7 @@ expect    |  expect.toPass @ a.test.ts:11
 test.step |    step 2, attempt: 0 @ a.test.ts:7
 test.step |    ↪ error: Error: expect(received).toBe(expected) // Object.is equality
 expect    |      expect.toBe @ a.test.ts:9
-expect    |        ↪ error: Error: expect(received).toBe(expected) // Object.is equality
+expect    |      ↪ error: Error: expect(received).toBe(expected) // Object.is equality
 test.step |    step 2, attempt: 1 @ a.test.ts:7
 expect    |      expect.toBe @ a.test.ts:9
 test.step |  step 3 @ a.test.ts:12
@@ -895,7 +894,6 @@ hook      |After Hooks
 
 test('library API call inside expect.toPass', async ({ runInlineTest }) => {
   test.info().annotations.push({ type: 'issue', description: 'https://github.com/microsoft/playwright/issues/30322' });
-  test.fixme();
   const result = await runInlineTest({
     'reporter.ts': stepIndentReporter,
     'playwright.config.ts': `
@@ -940,4 +938,53 @@ hook      |After Hooks
 fixture   |  fixture: page
 fixture   |  fixture: context
 `);
-});
+});
+
+test('library API call inside expect.poll', async ({ runInlineTest }) => {
+  test.info().annotations.push({ type: 'issue', description: 'https://github.com/microsoft/playwright/issues/30322' });
+  const result = await runInlineTest({
+    'reporter.ts': stepIndentReporter,
+    'playwright.config.ts': `
+      module.exports = {
+        reporter: './reporter',
+      };
+    `,
+    'a.test.ts': `
+      import { test, expect } from '@playwright/test';
+      test('pass', async ({page}) => {
+        let counter = 0
+        const a = [];
+        await expect.poll(async () => {
+          await page.goto('about:blank');
+          await test.step('inner step attempt: ' + counter, async () => {
+            counter++;
+            expect(1).toBe(1);
+          });
+          a.push(1);
+          return a;
+        }).toHaveLength(2);
+      });
+    `
+  }, { reporter: '', workers: 1 });
+
+  expect(result.exitCode).toBe(0);
+  expect(stripAnsi(result.output)).toBe(`
+hook      |Before Hooks
+fixture   |  fixture: browser
+pw:api    |    browserType.launch
+fixture   |  fixture: context
+pw:api    |    browser.newContext
+fixture   |  fixture: page
+pw:api    |    browserContext.newPage
+expect    |expect.poll.toHaveLength @ a.test.ts:14
+pw:api    |  page.goto(about:blank) @ a.test.ts:7
+test.step |  inner step attempt: 0 @ a.test.ts:8
+expect    |    expect.toBe @ a.test.ts:10
+pw:api    |  page.goto(about:blank) @ a.test.ts:7
+test.step |  inner step attempt: 1 @ a.test.ts:8
+expect    |    expect.toBe @ a.test.ts:10
+hook      |After Hooks
+fixture   |  fixture: page
+fixture   |  fixture: context
+`);
+});