fix(runner): test.extend doesn't work in hooks without test (#4065)

Author: Dunqing

packages/runner/src/fixture.ts

@@ -1,3 +1,4 @@
+import { getFixture } from './map'
 import type { TestContext } from './types'
 
 export interface FixtureItem {
@@ -43,32 +44,41 @@ export function mergeContextFixtures(fixtures: Record<string, any>, context: { f
   return context
 }
 
-export function withFixtures(fn: Function, fixtures: FixtureItem[], context: TestContext & Record<string, any>) {
-  if (!fixtures.length)
-    return () => fn(context)
+export function withFixtures(fn: Function, testContext?: TestContext) {
+  return (hookContext?: TestContext) => {
+    const context: TestContext & { [key: string]: any } | undefined = hookContext || testContext
 
-  const usedProps = getUsedProps(fn)
-  if (!usedProps.length)
-    return () => fn(context)
+    if (!context)
+      return fn({})
 
-  const usedFixtures = fixtures.filter(({ prop }) => usedProps.includes(prop))
-  const pendingFixtures = resolveDeps(usedFixtures)
-  let cursor = 0
+    const fixtures = getFixture(context)
+    if (!fixtures?.length)
+      return fn(context)
 
-  async function use(fixtureValue: any) {
-    const { prop } = pendingFixtures[cursor++]
-    context[prop] = fixtureValue
-    if (cursor < pendingFixtures.length)
-      await next()
-    else await fn(context)
-  }
+    const usedProps = getUsedProps(fn)
+    if (!usedProps.length)
+      return fn(context)
 
-  async function next() {
-    const { value } = pendingFixtures[cursor]
-    typeof value === 'function' ? await value(context, use) : await use(value)
-  }
+    const usedFixtures = fixtures.filter(({ prop }) => usedProps.includes(prop))
+    const pendingFixtures = resolveDeps(usedFixtures)
+    let cursor = 0
+
+    async function use(fixtureValue: any) {
+      const { prop } = pendingFixtures[cursor++]
+      context![prop] = fixtureValue
+
+      if (cursor < pendingFixtures.length)
+        await next()
+      else await fn(context)
+    }
+
+    async function next() {
+      const { value } = pendingFixtures[cursor]
+      typeof value === 'function' ? await value(context, use) : await use(value)
+    }
 
-  return () => next()
+    return next()
+  }
 }
 
 function resolveDeps(fixtures: FixtureItem[], depSet = new Set<FixtureItem>(), pendingFixtures: FixtureItem[] = []) {

packages/runner/src/hooks.ts

@@ -2,6 +2,7 @@ import type { OnTestFailedHandler, SuiteHooks, Test } from './types'
 import { getCurrentSuite, getRunner } from './suite'
 import { getCurrentTest } from './test-state'
 import { withTimeout } from './context'
+import { withFixtures } from './fixture'
 
 function getDefaultHookTimeout() {
   return getRunner().config.hookTimeout
@@ -15,10 +16,10 @@ export function afterAll(fn: SuiteHooks['afterAll'][0], timeout?: number) {
   return getCurrentSuite().on('afterAll', withTimeout(fn, timeout ?? getDefaultHookTimeout(), true))
 }
 export function beforeEach<ExtraContext = {}>(fn: SuiteHooks<ExtraContext>['beforeEach'][0], timeout?: number) {
-  return getCurrentSuite<ExtraContext>().on('beforeEach', withTimeout(fn, timeout ?? getDefaultHookTimeout(), true))
+  return getCurrentSuite<ExtraContext>().on('beforeEach', withTimeout(withFixtures(fn), timeout ?? getDefaultHookTimeout(), true))
 }
 export function afterEach<ExtraContext = {}>(fn: SuiteHooks<ExtraContext>['afterEach'][0], timeout?: number) {
-  return getCurrentSuite<ExtraContext>().on('afterEach', withTimeout(fn, timeout ?? getDefaultHookTimeout(), true))
+  return getCurrentSuite<ExtraContext>().on('afterEach', withTimeout(withFixtures(fn), timeout ?? getDefaultHookTimeout(), true))
 }
 
 export const onTestFailed = createTestHook<OnTestFailedHandler>('onTestFailed', (test, handler) => {

packages/runner/src/map.ts

@@ -1,8 +1,10 @@
 import type { Awaitable } from '@vitest/utils'
-import type { Suite, SuiteHooks, Test } from './types'
+import type { Suite, SuiteHooks, Test, TestContext } from './types'
+import type { FixtureItem } from './fixture'
 
 // use WeakMap here to make the Test and Suite object serializable
 const fnMap = new WeakMap()
+const fixtureMap = new WeakMap()
 const hooksMap = new WeakMap()
 
 export function setFn(key: Test, fn: (() => Awaitable<void>)) {
@@ -13,6 +15,14 @@ export function getFn<Task = Test>(key: Task): (() => Awaitable<void>) {
   return fnMap.get(key as any)
 }
 
+export function setFixture(key: TestContext, fixture: FixtureItem[] | undefined) {
+  fixtureMap.set(key, fixture)
+}
+
+export function getFixture<Context = TestContext>(key: Context): FixtureItem[] {
+  return fixtureMap.get(key as any)
+}
+
 export function setHooks(key: Suite, hooks: SuiteHooks) {
   hooksMap.set(key, hooks)
 }

packages/runner/src/suite.ts

@@ -3,7 +3,7 @@ import type { File, Fixtures, RunMode, Suite, SuiteAPI, SuiteCollector, SuiteFac
 import type { VitestRunner } from './types/runner'
 import { createChainable } from './utils/chain'
 import { collectTask, collectorContext, createTestContext, runWithSuite, withTimeout } from './context'
-import { getHooks, setFn, setHooks } from './map'
+import { getHooks, setFixture, setFn, setHooks } from './map'
 import type { FixtureItem } from './fixture'
 import { mergeContextFixtures, withFixtures } from './fixture'
 
@@ -96,10 +96,9 @@ function createSuiteCollector(name: string, factory: SuiteFactory = () => { }, m
       enumerable: false,
     })
 
+    setFixture(context, this.fixtures)
     setFn(test, withTimeout(
-      this.fixtures
-        ? withFixtures(fn, this.fixtures, context)
-        : () => fn(context),
+      withFixtures(fn, context),
       options?.timeout ?? runner.config.testTimeout,
     ))
 

test/core/test/test-extend-with-top-level-hooks.test.ts

@@ -0,0 +1,29 @@
+import { afterEach, beforeEach, expect, test } from 'vitest'
+
+interface Fixture { foo: number }
+
+const test1 = test.extend<Fixture>({
+  foo: 1,
+})
+
+const test2 = test.extend<Fixture>({
+  foo: 2,
+})
+
+test1('the foo should be 1', ({ foo }) => {
+  expect(foo).toBe(1)
+})
+
+test2('the foo should be 2', ({ foo }) => {
+  expect(foo).toBe(2)
+})
+
+let nextFoo = 1
+beforeEach<Fixture>(({ foo }) => {
+  expect(foo).toBe(nextFoo)
+})
+
+afterEach<Fixture>(({ foo }) => {
+  expect(foo).toBe(nextFoo)
+  nextFoo++
+})