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

hi-ogawa · web-flow · commit 0c2924b7aeb7 · 2025-03-03T17:11:37.000+01:00
diff --git a/packages/runner/src/hooks.ts b/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 },
+    ),
   )
 }
 
diff --git a/packages/runner/src/run.ts b/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))
     }
   }
 
diff --git a/test/cli/fixtures/fails/hooks-timeout-before-hook-cleanup-callback.test.ts b/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", () => {})
+})
diff --git a/test/cli/test/__snapshots__/fails.test.ts.snap b/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

Original file line number	Diff line number	Diff line change
`@@ -21,6 +21,7 @@ import { processError } from '@vitest/utils/error'`
`21`	`21`	`import { collectTests } from './collect'`
`22`	`22`	`import { PendingError } from './errors'`
`23`	`23`	`import { callFixtureCleanup } from './fixture'`
	`24`	`+import { getBeforeHookCleanupCallback } from './hooks'`
`24`	`25`	`import { getFn, getHooks } from './map'`
`25`	`26`	`import { setCurrentTest } from './test-state'`
`26`	`27`	`import { limitConcurrency } from './utils/limit-concurrency'`
`@@ -143,14 +144,18 @@ export async function callSuiteHook<T extends keyof SuiteHooks>(`
`143`	`144`	`updateSuiteHookState(currentTask, name, 'run', runner)`
`144`	`145`	`}`
`145`	`146`
	`147`	`+ async function runHook(hook: Function) {`
	`148`	`+ return getBeforeHookCleanupCallback(hook, await hook(...args))`
	`149`	`+ }`
	`150`	`+`
`146`	`151`	`if (sequence === 'parallel') {`
`147`	`152`	`callbacks.push(`
`148`		`- ...(await Promise.all(hooks.map(hook => (hook as any)(...args)))),`
	`153`	`+ ...(await Promise.all(hooks.map(hook => runHook(hook)))),`
`149`	`154`	`)`
`150`	`155`	`}`
`151`	`156`	`else {`
`152`	`157`	`for (const hook of hooks) {`
`153`		`- callbacks.push(await (hook as any)(...args))`
	`158`	`+ callbacks.push(await runHook(hook))`
`154`	`159`	`}`
`155`	`160`	`}`
`156`	`161`