From d3836e9c33eabfab1ed603e70c88d408ca6d4685 Mon Sep 17 00:00:00 2001
From: Adam Gastineau <adam@agg.im>
Date: Thu, 13 Feb 2025 11:13:28 -0800
Subject: [PATCH] Working warning annotation on expect

---
 packages/playwright/src/matchers/expect.ts   | 18 ++++++++++++++++--
 packages/playwright/src/worker/testInfo.ts   |  1 +
 packages/playwright/src/worker/workerMain.ts |  3 +++
 3 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/packages/playwright/src/matchers/expect.ts b/packages/playwright/src/matchers/expect.ts
index 707e5457e1..48a0b801d8 100644
--- a/packages/playwright/src/matchers/expect.ts
+++ b/packages/playwright/src/matchers/expect.ts
@@ -381,8 +381,22 @@ class ExpectMetaInfoProxyHandler implements ProxyHandler<any> {
         setMatcherCallContext({ expectInfo: this._info, testInfo, step: step.info });
         const callback = () => matcher.call(target, ...args);
         const result = zones.run('stepZone', step, callback);
-        if (result instanceof Promise)
-          return result.then(finalizer).catch(reportStepError);
+        if (result instanceof Promise) {
+          const promise = result.then(finalizer).catch(reportStepError);
+
+          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;
       } catch (e) {
diff --git a/packages/playwright/src/worker/testInfo.ts b/packages/playwright/src/worker/testInfo.ts
index 5e964af5e0..9ba4a7f7b2 100644
--- a/packages/playwright/src/worker/testInfo.ts
+++ b/packages/playwright/src/worker/testInfo.ts
@@ -99,6 +99,7 @@ export class TestInfoImpl implements TestInfo {
   duration: number = 0;
   readonly annotations: Annotation[] = [];
   readonly attachments: TestInfo['attachments'] = [];
+  readonly unusedAsyncApiCalls: Set<Promise<any>> = new Set();
   status: TestStatus = 'passed';
   snapshotSuffix: string = '';
   readonly outputDir: string;
diff --git a/packages/playwright/src/worker/workerMain.ts b/packages/playwright/src/worker/workerMain.ts
index e09579d0e0..bbe366a5bb 100644
--- a/packages/playwright/src/worker/workerMain.ts
+++ b/packages/playwright/src/worker/workerMain.ts
@@ -369,6 +369,9 @@ export class WorkerMain extends ProcessRunner {
         // Now run the test itself.
         const fn = test.fn; // Extract a variable to get a better stack trace ("myTest" vs "TestCase.myTest [as fn]").
         await fn(testFunctionParams, testInfo);
+        // Create warning if any of the async calls were not awaited.
+        if (testInfo.unusedAsyncApiCalls.size > 0)
+          testInfo.annotations.push({ type: 'warning', description: 'Some async calls were not awaited by the end of the test. This can cause flakiness.' });
       });
     }).catch(() => {});  // Ignore the top-level error, it is already inside TestInfo.errors.