fix: allow catch/finally for async assertion (#9827)

Author: hi-ogawa

packages/browser/src/client/utils.ts

@@ -54,9 +54,11 @@ export function ensureAwaited<T>(promise: (error?: Error) => Promise<T>): Promis
       return (promiseResult ||= promise(sourceError)).then(onFulfilled, onRejected)
     },
     catch(onRejected) {
+      awaited = true
       return (promiseResult ||= promise(sourceError)).catch(onRejected)
     },
     finally(onFinally) {
+      awaited = true
       return (promiseResult ||= promise(sourceError)).finally(onFinally)
     },
     [Symbol.toStringTag]: 'Promise',

packages/expect/src/utils.ts

@@ -71,9 +71,11 @@ export function recordAsyncExpect(
         return promise.then(onFulfilled, onRejected)
       },
       catch(onRejected) {
+        resolved = true
         return promise.catch(onRejected)
       },
       finally(onFinally) {
+        resolved = true
         return promise.finally(onFinally)
       },
       [Symbol.toStringTag]: 'Promise',

packages/vitest/src/integrations/chai/poll.ts

@@ -152,9 +152,11 @@ export function createExpectPoll(expect: ExpectStatic): ExpectStatic['poll'] {
               return (resultPromise ||= promise()).then(onFulfilled, onRejected)
             },
             catch(onRejected) {
+              awaited = true
               return (resultPromise ||= promise()).catch(onRejected)
             },
             finally(onFinally) {
+              awaited = true
               return (resultPromise ||= promise()).finally(onFinally)
             },
             [Symbol.toStringTag]: 'Promise',

test/browser/fixtures/failing/failing.test.ts

@@ -38,3 +38,15 @@ it('several locator methods are not awaited', () => {
 it('correctly prints error from a bundled file', () => {
   index()
 })
+
+it('not awaited but with then/catch/finally', async () => {
+  await page.getByRole('button').click().then()
+  await page.getByRole('button').click().catch()
+  await page.getByRole('button').click().finally()
+  await userEvent.click(page.getByRole('button')).then()
+  await userEvent.click(page.getByRole('button')).catch()
+  await userEvent.click(page.getByRole('button')).finally()
+  page.getByRole('button').click().then()
+  page.getByRole('button').click().catch()
+  page.getByRole('button').click().finally()
+})

test/browser/specs/runner.test.ts

@@ -209,7 +209,7 @@ error with a stack
 })
 
 test(`stack trace points to correct file in every browser when failed`, async () => {
-  expect.assertions(29)
+  expect.assertions(30)
   const { stderr } = await runBrowserTests({
     root: './fixtures/failing',
     reporters: [
@@ -273,6 +273,9 @@ test(`stack trace points to correct file in every browser when failed`, async ()
 
   // index() is called from a bundled file
   expect(stderr).toMatch(/failing.test.ts:39:(2|8)/)
+
+  // "not awaited but with then/catch/finally" test should not produce warnings
+  expect(stderr).not.toMatch(/failing.test.ts:4[3-8]/)
 })
 
 test('user-event', async () => {

test/cli/test/fails.test.ts

@@ -117,6 +117,119 @@ it('prints a warning if the assertion is not awaited', async () => {
   `)
 })
 
+it('no async tracking after then/catch/finally', async () => {
+  // resolve + then/catch/finally
+  // rejects + then/catch/finally
+  // expect.poll + then/catch/finally
+  const { stderr, errorTree } = await runInlineTests({
+    'await.test.js': ts`
+      import { expect, test } from 'vitest';
+
+      test('resolves + then', async () => {
+        await expect(Promise.resolve(1)).resolves.toBe(1).then()
+      })
+
+      test('resolves + catch', async () => {
+        await expect(Promise.resolve(1)).resolves.toBe(1).catch()
+      })
+
+      test('resolves + finally', async () => {
+        await expect(Promise.resolve(1)).resolves.toBe(1).finally()
+      })
+
+      test('rejects + then', async () => {
+        await expect(Promise.reject(1)).rejects.toBe(1).then()
+      })
+
+      test('rejects + catch', async () => {
+        await expect(Promise.reject(1)).rejects.toBe(1).catch()
+      })
+
+      test('rejects + finally', async () => {
+        await expect(Promise.reject(1)).rejects.toBe(1).finally()
+      })
+
+      test('expect.poll + then', async () => {
+        await expect.poll(() => 2).toBe(2).then()
+      })
+
+      test('expect.poll + catch', async () => {
+        await expect.poll(() => 2).toBe(2).catch()
+      })
+
+      test('expect.poll + finally', async () => {
+        await expect.poll(() => 2).toBe(2).finally()
+      })
+    `,
+    'hang.test.js': ts`
+      import { expect, test } from 'vitest';
+
+      test('resolves + then', async () => {
+        expect(Promise.resolve(1)).resolves.toBe(1).then()
+      })
+
+      test('resolves + catch', async () => {
+        expect(Promise.resolve(1)).resolves.toBe(1).catch()
+      })
+
+      test('resolves + finally', async () => {
+        expect(Promise.resolve(1)).resolves.toBe(1).finally()
+      })
+
+      test('rejects + then', async () => {
+        expect(Promise.reject(1)).rejects.toBe(1).then()
+      })
+
+      test('rejects + catch', async () => {
+        expect(Promise.reject(1)).rejects.toBe(1).catch()
+      })
+
+      test('rejects + finally', async () => {
+        expect(Promise.reject(1)).rejects.toBe(1).finally()
+      })
+
+      test('expect.poll + then', async () => {
+        expect.poll(() => 2).toBe(2).then()
+      })
+
+      test('expect.poll + catch', async () => {
+        expect.poll(() => 2).toBe(2).catch()
+      })
+
+      test('expect.poll + finally', async () => {
+        expect.poll(() => 2).toBe(2).finally()
+      })
+    `,
+  })
+  expect(stderr).toMatchInlineSnapshot(`""`)
+  expect(errorTree()).toMatchInlineSnapshot(`
+    {
+      "await.test.js": {
+        "expect.poll + catch": "passed",
+        "expect.poll + finally": "passed",
+        "expect.poll + then": "passed",
+        "rejects + catch": "passed",
+        "rejects + finally": "passed",
+        "rejects + then": "passed",
+        "resolves + catch": "passed",
+        "resolves + finally": "passed",
+        "resolves + then": "passed",
+      },
+      "hang.test.js": {
+        "expect.poll + catch": "passed",
+        "expect.poll + finally": "passed",
+        "expect.poll + then": "passed",
+        "rejects + catch": "passed",
+        "rejects + finally": "passed",
+        "rejects + then": "passed",
+        "resolves + catch": "passed",
+        "resolves + finally": "passed",
+        "resolves + then": "passed",
+      },
+    }
+  `)
+})
+
 it('prints a warning if the assertion is not awaited in the browser mode', async () => {
   const { stderr } = await runInlineTests({
     'base.test.js': ts`