From cc85b1e832e923bb8d3c8b894c43a5f7fb8d2471 Mon Sep 17 00:00:00 2001
From: Dunqing <dengqing0821@gmail.com>
Date: Sun, 24 Sep 2023 17:29:13 +0800
Subject: [PATCH] fix: incorrect execution order in nested describe

---
 packages/runner/src/fixture.ts     | 34 +++++++++----
 packages/runner/src/run.ts         |  2 +-
 test/core/test/test-extend.test.ts | 78 ++++++++++++++++++++++++++++++
 3 files changed, 104 insertions(+), 10 deletions(-)

diff --git a/packages/runner/src/fixture.ts b/packages/runner/src/fixture.ts
index ea8791952e03..9fc94850d105 100644
--- a/packages/runner/src/fixture.ts
+++ b/packages/runner/src/fixture.ts
@@ -45,13 +45,17 @@ export function mergeContextFixtures(fixtures: Record<string, any>, context: { f
 }
 
 const fixtureValueMap = new Map<FixtureItem, any>()
-const fixtureCleanupSet = new Set<() => void | Promise<void>>()
+const fixtureCleanupFnMap = new Map<string, Array<() => void | Promise<void>>>()
 
-export async function callFixtureCleanup() {
-  await Promise.all([...fixtureCleanupSet].map(async (fn) => {
-    await fn()
-  }))
-  fixtureValueMap.clear()
+export async function callFixtureCleanup(id: string) {
+  const cleanupFnArray = fixtureCleanupFnMap.get(id)
+  if (!cleanupFnArray)
+    return
+
+  for (const cleanup of cleanupFnArray.reverse())
+    await cleanup()
+
+  fixtureCleanupFnMap.delete(id)
 }
 
 export function withFixtures(fn: Function, testContext?: TestContext) {
@@ -61,6 +65,12 @@ export function withFixtures(fn: Function, testContext?: TestContext) {
     if (!context)
       return fn({})
 
+    let cleanupFnArray = fixtureCleanupFnMap.get(context.task.suite.id)!
+    if (!cleanupFnArray) {
+      cleanupFnArray = []
+      fixtureCleanupFnMap.set(context.task.suite.id, cleanupFnArray)
+    }
+
     const fixtures = getFixture(context)
     if (!fixtures?.length)
       return fn(context)
@@ -77,7 +87,13 @@ export function withFixtures(fn: Function, testContext?: TestContext) {
       async function use(fixtureValue: any) {
         const fixture = pendingFixtures[cursor++]
         context![fixture.prop] = fixtureValue
-        fixtureValueMap.set(fixture, fixtureValue)
+
+        if (!fixtureValueMap.has(fixture)) {
+          fixtureValueMap.set(fixture, fixtureValue)
+          cleanupFnArray.unshift(() => {
+            fixtureValueMap.delete(fixture)
+          })
+        }
 
         if (cursor < pendingFixtures.length) {
           await next()
@@ -91,7 +107,7 @@ export function withFixtures(fn: Function, testContext?: TestContext) {
             reject(err)
           }
           return new Promise<void>((resolve) => {
-            fixtureCleanupSet.add(resolve)
+            cleanupFnArray.push(resolve)
           })
         }
       }
@@ -106,7 +122,7 @@ export function withFixtures(fn: Function, testContext?: TestContext) {
       }
 
       const setupFixturePromise = next()
-      fixtureCleanupSet.add(() => setupFixturePromise)
+      cleanupFnArray.unshift(() => setupFixturePromise)
     })
   }
 }
diff --git a/packages/runner/src/run.ts b/packages/runner/src/run.ts
index 50c209242419..798ac93f9633 100644
--- a/packages/runner/src/run.ts
+++ b/packages/runner/src/run.ts
@@ -322,7 +322,7 @@ export async function runSuite(suite: Suite, runner: VitestRunner) {
     }
 
     try {
-      await callFixtureCleanup()
+      await callFixtureCleanup(suite.id)
       await callSuiteHook(suite, suite, 'afterAll', runner, [suite])
       await callCleanupHooks(beforeAllCleanups)
     }
diff --git a/test/core/test/test-extend.test.ts b/test/core/test/test-extend.test.ts
index 21c970122958..6da9ac6b2851 100644
--- a/test/core/test/test-extend.test.ts
+++ b/test/core/test/test-extend.test.ts
@@ -199,4 +199,82 @@ describe('test.extend()', () => {
       expect(teardownFn).toBeCalledTimes(2)
     })
   })
+
+  describe('fixture in nested describe', () => {
+    interface Fixture {
+      foo: number
+      bar: number
+    }
+
+    const fooFn = vi.fn(() => 0)
+    const fooCleanup = vi.fn()
+
+    const barFn = vi.fn(() => 0)
+    const barCleanup = vi.fn()
+
+    const nestedTest = test.extend<Fixture>({
+      async foo({}, use) {
+        await use(fooFn())
+        fooCleanup()
+      },
+      async bar({}, use) {
+        await use(barFn())
+        barCleanup()
+      },
+    })
+
+    beforeEach<Fixture>(({ foo }) => {
+      expect(foo).toBe(0)
+    })
+
+    nestedTest('should only initialize foo', ({ foo }) => {
+      expect(foo).toBe(0)
+      expect(fooFn).toBeCalledTimes(1)
+      expect(barFn).toBeCalledTimes(0)
+    })
+
+    describe('level 2, using both foo and bar together', () => {
+      beforeEach<Fixture>(({ foo, bar }) => {
+        expect(foo).toBe(0)
+        expect(bar).toBe(0)
+      })
+
+      nestedTest('should only initialize bar', ({ foo, bar }) => {
+        expect(foo).toBe(0)
+        expect(bar).toBe(0)
+        expect(fooFn).toBeCalledTimes(1)
+        expect(barFn).toBeCalledTimes(1)
+      })
+
+      afterEach<Fixture>(({ foo, bar }) => {
+        expect(foo).toBe(0)
+        expect(bar).toBe(0)
+      })
+
+      afterAll(() => {
+        // foo setup in outside describe
+        // cleanup also called in outside describe
+        expect(fooCleanup).toHaveBeenCalledTimes(0)
+        // bar setup in inside describe
+        // cleanup also called in inside describe
+        expect(barCleanup).toHaveBeenCalledTimes(1)
+      })
+    })
+
+    nestedTest('level 2 will not call foo cleanup', ({ foo }) => {
+      expect(foo).toBe(0)
+      expect(fooFn).toBeCalledTimes(1)
+    })
+
+    afterEach<Fixture>(({ foo }) => {
+      expect(foo).toBe(0)
+    })
+
+    afterAll(() => {
+      // foo setup in this describe
+      // cleanup also called in this describe
+      expect(fooCleanup).toHaveBeenCalledTimes(1)
+      expect(barCleanup).toHaveBeenCalledTimes(1)
+    })
+  })
 })