Generalized implementation and step API

packages/playwright/src/common/testType.ts

@@ -25,6 +25,7 @@ import { wrapFunctionWithLocation } from '../transform/transform';
 import type { FixturesWithLocation } from './config';
 import type { Fixtures, TestDetails, TestStepInfo, TestType } from '../../types/test';
 import type { Location } from '../../types/testReporter';
+import type { TestInfoImpl } from '../worker/testInfo';
 
 
 const testTypeSymbol = Symbol('testType');
@@ -261,10 +262,14 @@ export class TestTypeImpl {
     suite._use.push({ fixtures, location });
   }
 
-  async _step<T>(expectation: 'pass'|'skip', title: string, body: (step: TestStepInfo) => T | Promise<T>, options: {box?: boolean, location?: Location, timeout?: number } = {}): Promise<T> {
+  _step<T>(expectation: 'pass'|'skip', title: string, body: (step: TestStepInfo) => T | Promise<T>, options: {box?: boolean, location?: Location, timeout?: number } = {}): Promise<T> {
     const testInfo = currentTestInfo();
     if (!testInfo)
       throw new Error(`test.step() can only be called from a test`);
+    return testInfo._wrapPromiseAPIResult(this._stepInternal(expectation, testInfo, title, body, options));
+  }
+
+  async _stepInternal<T>(expectation: 'pass'|'skip', testInfo: TestInfoImpl, title: string, body: (step: TestStepInfo) => T | Promise<T>, options: {box?: boolean, location?: Location, timeout?: number } = {}): Promise<T> {
     const step = testInfo._addStep({ category: 'test.step', title, location: options.location, box: options.box });
     return await zones.run('stepZone', step, async () => {
       try {

packages/playwright/src/matchers/expect.ts

@@ -383,19 +383,7 @@ class ExpectMetaInfoProxyHandler implements ProxyHandler<any> {
         const result = zones.run('stepZone', step, callback);
         if (result instanceof Promise) {
           const promise = result.then(finalizer).catch(reportStepError);
-
+          return testInfo._wrapPromiseAPIResult(promise);
-          testInfo?.unusedAsyncApiCalls.add(promise);
-
-          const oldThen = promise.then;
-          promise.then = ((...args: any[]) => {
-            if (args[0] !== undefined) {
-              // onfulfilled callback
-              testInfo?.unusedAsyncApiCalls.delete(promise);
-            }
-            return oldThen.call(promise, ...args);
-          }) as any;
-
-          return promise;
         }
         finalizer();
         return result;

packages/playwright/src/worker/testInfo.ts

@@ -411,6 +411,26 @@ export class TestInfoImpl implements TestInfo {
     this._timeoutManager.setIgnoreTimeouts();
   }
 
+  /**
+   * Enables a promise API call to be tracked by the test, alerting if unawaited.
+   *
+   * **NOTE:** Returning from an async function wraps the result in a promise, regardless of whether the return value is a promise. This will automatically mark the promise as awaited. Avoid this.
+   */
+  _wrapPromiseAPIResult<T>(promise: Promise<T>): Promise<T> {
+    this.unusedAsyncApiCalls.add(promise);
+
+    const oldThen = promise.then;
+    promise.then = ((...args: any[]) => {
+      if (args[0] !== undefined) {
+        // onfulfilled callback, which means .then() was called
+        this.unusedAsyncApiCalls.delete(promise);
+      }
+      return oldThen.call(promise, ...args);
+    }) as any;
+
+    return promise;
+  }
+
   // ------------ TestInfo methods ------------
 
   async attach(name: string, options: { path?: string, body?: string | Buffer, contentType?: string } = {}) {