From e7c5e1f7b65526d268b905e2da9001bd61d96f0d Mon Sep 17 00:00:00 2001
From: Hiroshi Ogawa <hi.ogawa.zz@gmail.com>
Date: Thu, 7 Dec 2023 19:08:20 +0900
Subject: [PATCH] fix(runner): fix fixture cleanup when test times out (#4679)

---
 packages/runner/src/fixture.ts                | 65 ++++++++-----------
 test/core/test/test-extend.test.ts            | 24 +++++++
 .../test-extend/fixture-error.test.ts         |  3 +-
 .../test/__snapshots__/runner.test.ts.snap    |  3 +-
 4 files changed, 54 insertions(+), 41 deletions(-)

diff --git a/packages/runner/src/fixture.ts b/packages/runner/src/fixture.ts
index 026cdf4f19ab..c210b09df1e7 100644
--- a/packages/runner/src/fixture.ts
+++ b/packages/runner/src/fixture.ts
@@ -1,3 +1,4 @@
+import { createDefer } from '@vitest/utils'
 import { getFixture } from './map'
 import type { TestContext } from './types'
 
@@ -78,49 +79,37 @@ export function withFixtures(fn: Function, testContext?: TestContext) {
     if (!pendingFixtures.length)
       return fn(context)
 
-    let cursor = 0
-
-    return new Promise((resolve, reject) => {
-      async function use(fixtureValue: any) {
-        const fixture = pendingFixtures[cursor++]
-        context![fixture.prop] = fixtureValue
-
-        if (!fixtureValueMap.has(fixture)) {
-          fixtureValueMap.set(fixture, fixtureValue)
-          cleanupFnArray.unshift(() => {
-            fixtureValueMap.delete(fixture)
-          })
-        }
-
-        if (cursor < pendingFixtures.length) {
-          await next()
+    async function resolveFixtures() {
+      for (const fixture of pendingFixtures) {
+        // fixture could be already initialized during "before" hook
+        if (fixtureValueMap.has(fixture))
+          continue
+
+        let resolvedValue: unknown
+        if (fixture.isFn) {
+          // wait for `use` call to extract fixture value
+          const useFnArgPromise = createDefer()
+          fixture.value(context, async (useFnArg: unknown) => {
+            useFnArgPromise.resolve(useFnArg)
+            // suspend fixture teardown until cleanup
+            const teardownPromise = createDefer<void>()
+            cleanupFnArray.push(teardownPromise.resolve)
+            await teardownPromise
+          }).catch(useFnArgPromise.reject) // treat fixture function error (until `use` call) as test failure
+          resolvedValue = await useFnArgPromise
         }
         else {
-          // When all fixtures setup, call the test function
-          try {
-            resolve(await fn(context))
-          }
-          catch (err) {
-            reject(err)
-          }
-          return new Promise<void>((resolve) => {
-            cleanupFnArray.push(resolve)
-          })
+          resolvedValue = fixture.value
         }
+        context![fixture.prop] = resolvedValue
+        fixtureValueMap.set(fixture, resolvedValue)
+        cleanupFnArray.unshift(() => {
+          fixtureValueMap.delete(fixture)
+        })
       }
+    }
 
-      async function next() {
-        const fixture = pendingFixtures[cursor]
-        const { isFn, value } = fixture
-        if (fixtureValueMap.has(fixture))
-          return use(fixtureValueMap.get(fixture))
-        else
-          return isFn ? value(context, use) : use(value)
-      }
-
-      const setupFixturePromise = next().catch(reject)
-      cleanupFnArray.unshift(() => setupFixturePromise)
-    })
+    return resolveFixtures().then(() => fn(context))
   }
 }
 
diff --git a/test/core/test/test-extend.test.ts b/test/core/test/test-extend.test.ts
index de191aa10911..3e2e4a24ba75 100644
--- a/test/core/test/test-extend.test.ts
+++ b/test/core/test/test-extend.test.ts
@@ -177,6 +177,28 @@ describe('test.extend()', () => {
     })
   })
 
+  describe('fixture only in beforeEach', () => {
+    beforeEach<Fixtures>(({ todoList }) => {
+      expect(todoList).toEqual([1, 2, 3])
+      expect(todoFn).toBeCalledTimes(1)
+    })
+
+    myTest('no fixture in test', () => {
+      expect(todoFn).toBeCalledTimes(1)
+    })
+  })
+
+  describe('fixture only in afterEach', () => {
+    afterEach<Fixtures>(({ todoList }) => {
+      expect(todoList).toEqual([1, 2, 3])
+      expect(todoFn).toBeCalledTimes(1)
+    })
+
+    myTest('no fixture in test', () => {
+      expect(todoFn).toBeCalledTimes(0)
+    })
+  })
+
   describe('fixture call times', () => {
     const apiFn = vi.fn(() => true)
     const serviceFn = vi.fn(() => true)
@@ -203,6 +225,8 @@ describe('test.extend()', () => {
     beforeEach<APIFixture>(({ api, service }) => {
       expect(api).toBe(true)
       expect(service).toBe(true)
+      expect(apiFn).toBeCalledTimes(1)
+      expect(serviceFn).toBeCalledTimes(1)
     })
 
     testAPI('Should init 1 time', ({ api }) => {
diff --git a/test/fails/fixtures/test-extend/fixture-error.test.ts b/test/fails/fixtures/test-extend/fixture-error.test.ts
index f02431a627b4..4f815b9cae8d 100644
--- a/test/fails/fixtures/test-extend/fixture-error.test.ts
+++ b/test/fails/fixtures/test-extend/fixture-error.test.ts
@@ -39,8 +39,7 @@ describe('error thrown in test fixtures', () => {
   myTest('fixture errors', ({ a }) => {})
 })
 
-// TODO: enable when #4669 is fixed
-describe.skip('correctly fails when test times out', () => {
+describe('correctly fails when test times out', () => {
   const myTest = test.extend<{ a: number }>({
     a: async ({}, use) => {
       await use(2)
diff --git a/test/fails/test/__snapshots__/runner.test.ts.snap b/test/fails/test/__snapshots__/runner.test.ts.snap
index 23b1824e9134..b628660bd975 100644
--- a/test/fails/test/__snapshots__/runner.test.ts.snap
+++ b/test/fails/test/__snapshots__/runner.test.ts.snap
@@ -43,7 +43,8 @@ TypeError: failure"
 exports[`should fail test-extend/circular-dependency.test.ts > test-extend/circular-dependency.test.ts 1`] = `"Error: Circular fixture dependency detected: a <- b <- a"`;
 
 exports[`should fail test-extend/fixture-error.test.ts > test-extend/fixture-error.test.ts 1`] = `
-"Error: Error thrown in test fixture
+"Error: Test timed out in 20ms.
+Error: Error thrown in test fixture
 Error: Error thrown in afterEach fixture
 Error: Error thrown in beforeEach fixture"
 `;