fix(runner): fix beforeEach/All cleanup callback timeout

Author: hi-ogawa

packages/runner/src/hooks.ts

@@ -14,10 +14,22 @@ import { withFixtures } from './fixture'
 import { getCurrentSuite, getRunner } from './suite'
 import { getCurrentTest } from './test-state'
 
-function getDefaultHookTimeout() {
+function getDefaultHookTimeout(): number {
   return getRunner().config.hookTimeout
 }
 
+const CLEANUP_TIMEOUT_KEY = Symbol.for('VITEST_CLEANUP_TIMEOUT')
+
+export function getBeforeHookCleanupCallback(hook: Function, result: any): Function | undefined {
+  if (typeof result === 'function') {
+    const timeout
+      = CLEANUP_TIMEOUT_KEY in hook && typeof hook[CLEANUP_TIMEOUT_KEY] === 'number'
+        ? hook[CLEANUP_TIMEOUT_KEY]
+        : getDefaultHookTimeout()
+    return withTimeout(result, timeout, true)
+  }
+}
+
 /**
  * Registers a callback function to be executed once before all tests within the current suite.
  * This hook is useful for scenarios where you need to perform setup operations that are common to all tests in a suite, such as initializing a database connection or setting up a test environment.
@@ -35,11 +47,14 @@ function getDefaultHookTimeout() {
  * });
  * ```
  */
-export function beforeAll(fn: BeforeAllListener, timeout?: number): void {
+export function beforeAll(
+  fn: BeforeAllListener,
+  timeout: number = getDefaultHookTimeout(),
+): void {
   assertTypes(fn, '"beforeAll" callback', ['function'])
   return getCurrentSuite().on(
     'beforeAll',
-    withTimeout(fn, timeout ?? getDefaultHookTimeout(), true),
+    Object.assign(withTimeout(fn, timeout, true), { [CLEANUP_TIMEOUT_KEY]: timeout }),
   )
 }
 
@@ -87,12 +102,15 @@ export function afterAll(fn: AfterAllListener, timeout?: number): void {
  */
 export function beforeEach<ExtraContext = object>(
   fn: BeforeEachListener<ExtraContext>,
-  timeout?: number,
+  timeout: number = getDefaultHookTimeout(),
 ): void {
   assertTypes(fn, '"beforeEach" callback', ['function'])
   return getCurrentSuite<ExtraContext>().on(
     'beforeEach',
-    withTimeout(withFixtures(fn), timeout ?? getDefaultHookTimeout(), true),
+    Object.assign(
+      withTimeout(withFixtures(fn), timeout ?? getDefaultHookTimeout(), true),
+      { [CLEANUP_TIMEOUT_KEY]: timeout },
+    ),
   )
 }
 

packages/runner/src/run.ts

@@ -21,6 +21,7 @@ import { processError } from '@vitest/utils/error'
 import { collectTests } from './collect'
 import { PendingError } from './errors'
 import { callFixtureCleanup } from './fixture'
+import { getBeforeHookCleanupCallback } from './hooks'
 import { getFn, getHooks } from './map'
 import { setCurrentTest } from './test-state'
 import { limitConcurrency } from './utils/limit-concurrency'
@@ -143,14 +144,18 @@ export async function callSuiteHook<T extends keyof SuiteHooks>(
     updateSuiteHookState(currentTask, name, 'run', runner)
   }
 
+  async function runHook(hook: Function) {
+    return getBeforeHookCleanupCallback(hook, await hook(...args))
+  }
+
   if (sequence === 'parallel') {
     callbacks.push(
-      ...(await Promise.all(hooks.map(hook => (hook as any)(...args)))),
+      ...(await Promise.all(hooks.map(hook => runHook(hook)))),
     )
   }
   else {
     for (const hook of hooks) {
-      callbacks.push(await (hook as any)(...args))
+      callbacks.push(await runHook(hook))
     }
   }