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

Author: hi-ogawa

packages/runner/src/hooks.ts

@@ -18,6 +18,18 @@ function getDefaultHookTimeout() {
   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))
     }
   }
 

test/cli/fixtures/fails/hooks-timeout-before-hook-cleanup-callback.test.ts

@@ -0,0 +1,11 @@
+import { beforeEach, beforeAll, describe, test } from 'vitest';
+
+describe('beforeEach cleanup timeout', () => {
+  beforeEach(() => new Promise(() => {}), 101)
+  test("ok", () => {})
+})
+
+describe('beforeAll cleanup timeout', () => {
+  beforeAll(() => new Promise(() => {}), 102)
+  test("ok", () => {})
+})

test/cli/test/__snapshots__/fails.test.ts.snap

@@ -42,6 +42,11 @@ exports[`should fail hooks-fail-beforeAll.test.ts 1`] = `"TypeError: "beforeAll"
 
 exports[`should fail hooks-fail-beforeEach.test.ts 1`] = `"TypeError: "beforeEach" callback value must be function, received "string""`;
 
+exports[`should fail hooks-timeout-before-hook-cleanup-callback.test.ts 1`] = `
+"Error: Hook timed out in 101ms.
+Error: Hook timed out in 102ms."
+`;
+
 exports[`should fail inline-snapshop-inside-each.test.ts 1`] = `
 "Error: InlineSnapshot cannot be used inside of test.each or describe.each
 Error: InlineSnapshot cannot be used inside of test.each or describe.each