feat: add hooks with type-safe extra context to TestAPI (#8623)

Author: ysfaran

docs/guide/test-context.md

@@ -455,4 +455,26 @@ test('types are correct', ({
   // ...
 })
 ```
+
 :::
+
+When using `test.extend`, the extended `test` object provides type-safe `beforeEach` and `afterEach` hooks that are aware of the new context:
+
+```ts
+const test = baseTest.extend<{
+  todos: number[]
+}>({
+  todos: async ({}, use) => {
+    await use([])
+  },
+})
+
+// Unlike global hooks, these hooks are aware of the extended context
+test.beforeEach(({ todos }) => {
+  todos.push(1)
+})
+
+test.afterEach(({ todos }) => {
+  console.log(todos)
+})
+```

packages/runner/src/suite.ts

@@ -33,6 +33,7 @@ import {
   withTimeout,
 } from './context'
 import { mergeContextFixtures, mergeScopedFixtures, withFixtures } from './fixture'
+import { afterAll, afterEach, beforeAll, beforeEach } from './hooks'
 import { getHooks, setFn, setHooks, setTestFixture } from './map'
 import { getCurrentTest } from './test-state'
 import { findTestFileStackTrace } from './utils'
@@ -794,6 +795,11 @@ export function createTaskCollector(
     }, _context)
   }
 
+  taskFn.beforeEach = beforeEach
+  taskFn.afterEach = afterEach
+  taskFn.beforeAll = beforeAll
+  taskFn.afterAll = afterAll
+
   const _test = createChainable(
     ['concurrent', 'sequential', 'skip', 'only', 'todo', 'fails'],
     taskFn,

packages/runner/src/types/tasks.ts

@@ -1,5 +1,6 @@
 import type { Awaitable, TestError } from '@vitest/utils'
 import type { FixtureItem } from '../fixture'
+import type { afterAll, afterEach, beforeAll, beforeEach } from '../hooks'
 import type { ChainableFunction } from '../utils/chain'
 
 export type RunMode = 'run' | 'skip' | 'only' | 'todo' | 'queued'
@@ -482,8 +483,15 @@ interface ExtendedAPI<ExtraContext> {
   runIf: (condition: any) => ChainableTestAPI<ExtraContext>
 }
 
+interface Hooks<ExtraContext> {
+  beforeAll: typeof beforeAll
+  afterAll: typeof afterAll
+  beforeEach: typeof beforeEach<ExtraContext>
+  afterEach: typeof afterEach<ExtraContext>
+}
+
 export type TestAPI<ExtraContext = object> = ChainableTestAPI<ExtraContext>
-  & ExtendedAPI<ExtraContext> & {
+  & ExtendedAPI<ExtraContext> & Hooks<ExtraContext> & {
     extend: <T extends Record<string, any> = object>(
       fixtures: Fixtures<T, ExtraContext>
     ) => TestAPI<{

test/core/test/test-extend.test.ts

@@ -497,3 +497,33 @@ describe('suite with timeout', () => {
     expect(task.timeout).toBe(1_000)
   })
 }, 100)
+
+describe('type-safe fixture hooks', () => {
+  const counterTest = test.extend<{
+    counter: { value: number }
+    fileCounter: { value: number }
+  }>({
+    counter: async ({}, use) => { await use({ value: 0 }) },
+    fileCounter: [async ({}, use) => { await use({ value: 0 }) }, { scope: 'file' }],
+  })
+
+  counterTest.beforeEach(({ counter }) => {
+    // shouldn't have typescript error because of 'counter' here
+    counter.value += 1
+  })
+
+  counterTest.afterEach(({ fileCounter }) => {
+    // shouldn't have typescript error because of 'fileCounter' here
+    fileCounter.value += 2
+  })
+
+  // beforeAll and afterAll hooks are not tested here, because they don't provide an extra context
+
+  counterTest('beforeEach fixture hook can adapt type-safe context', ({ counter }) => {
+    expect(counter.value).toBe(1)
+  })
+
+  counterTest('afterEach fixture hook can adapt type-safe context', ({ fileCounter }) => {
+    expect(fileCounter.value).toBe(2)
+  })
+})