feat!: pass down context to test hooks (#7034)

Author: sheremet-va

Diff for: docs/api/index.md

@@ -1126,7 +1126,7 @@ These hooks will throw an error if they are called outside of the test body.
 
 ### onTestFinished {#ontestfinished}
 
-This hook is always called after the test has finished running. It is called after `afterEach` hooks since they can influence the test result. It receives a `TaskResult` object with the current test result.
+This hook is always called after the test has finished running. It is called after `afterEach` hooks since they can influence the test result. It receives an `ExtendedContext` object like `beforeEach` and `afterEach`.
 
 ```ts {1,5}
 import { onTestFinished, test } from 'vitest'
@@ -1183,7 +1183,7 @@ This hook is always called in reverse order and is not affected by [`sequence.ho
 
 ### onTestFailed
 
-This hook is called only after the test has failed. It is called after `afterEach` hooks since they can influence the test result. It receives a `TaskResult` object with the current test result. This hook is useful for debugging.
+This hook is called only after the test has failed. It is called after `afterEach` hooks since they can influence the test result. It receives an `ExtendedContext` object like `beforeEach` and `afterEach`. This hook is useful for debugging.
 
 ```ts {1,5-7}
 import { onTestFailed, test } from 'vitest'

Diff for: docs/guide/migration.md

@@ -42,6 +42,10 @@ import {
 
 If you are using `getCurrentSuite().custom()`, the `type` of the returned task is now is equal to `'test'`. The `Custom` type will be removed in Vitest 4.
 
+### `onTestFinished` and `onTestFailed` Now Receive a Context
+
+The [`onTestFinished`](/api/#ontestfinished) and [`onTestFailed`](/api/#ontestfailed) hooks previously received a test result as the first argument. Now, they receive a test context, like `beforeEach` and `afterEach`.
+
 ## Migrating to Vitest 2.0
 
 ### Default Pool is `forks`

Diff for: packages/runner/src/context.ts

@@ -71,14 +71,18 @@ export function createTestContext<T extends Test>(
     throw new PendingError('test is skipped; abort execution', test, note)
   }
 
-  context.onTestFailed = (fn) => {
+  context.onTestFailed = (handler, timeout) => {
     test.onFailed ||= []
-    test.onFailed.push(fn)
+    test.onFailed.push(
+      withTimeout(handler, timeout ?? runner.config.hookTimeout, true),
+    )
   }
 
-  context.onTestFinished = (fn) => {
+  context.onTestFinished = (handler, timeout) => {
     test.onFinished ||= []
-    test.onFinished.push(fn)
+    test.onFinished.push(
+      withTimeout(handler, timeout ?? runner.config.hookTimeout, true),
+    )
   }
 
   return (runner.extendTaskContext?.(context) as ExtendedContext<T>) || context

Diff for: packages/runner/src/run.ts

@@ -2,6 +2,7 @@ import type { Awaitable } from '@vitest/utils'
 import type { DiffOptions } from '@vitest/utils/diff'
 import type { FileSpec, VitestRunner } from './types/runner'
 import type {
+  ExtendedContext,
   File,
   HookCleanupCallback,
   HookListener,
@@ -62,32 +63,48 @@ function getSuiteHooks(
 
 async function callTestHooks(
   runner: VitestRunner,
-  task: Task,
-  hooks: ((result: TaskResult) => Awaitable<void>)[],
+  test: Test,
+  hooks: ((context: ExtendedContext<Test>) => Awaitable<void>)[],
   sequence: SequenceHooks,
 ) {
   if (sequence === 'stack') {
     hooks = hooks.slice().reverse()
   }
 
+  if (!hooks.length) {
+    return
+  }
+
+  const onTestFailed = test.context.onTestFailed
+  const onTestFinished = test.context.onTestFinished
+  test.context.onTestFailed = () => {
+    throw new Error(`Cannot call "onTestFailed" inside a test hook.`)
+  }
+  test.context.onTestFinished = () => {
+    throw new Error(`Cannot call "onTestFinished" inside a test hook.`)
+  }
+
   if (sequence === 'parallel') {
     try {
-      await Promise.all(hooks.map(fn => fn(task.result!)))
+      await Promise.all(hooks.map(fn => fn(test.context)))
     }
     catch (e) {
-      failTask(task.result!, e, runner.config.diffOptions)
+      failTask(test.result!, e, runner.config.diffOptions)
     }
   }
   else {
     for (const fn of hooks) {
       try {
-        await fn(task.result!)
+        await fn(test.context)
       }
       catch (e) {
-        failTask(task.result!, e, runner.config.diffOptions)
+        failTask(test.result!, e, runner.config.diffOptions)
       }
     }
   }
+
+  test.context.onTestFailed = onTestFailed
+  test.context.onTestFinished = onTestFinished
 }
 
 export async function callSuiteHook<T extends keyof SuiteHooks>(

Diff for: packages/runner/src/types/tasks.ts

@@ -608,12 +608,12 @@ export interface TaskContext<Task extends Test = Test> {
   /**
    * Extract hooks on test failed
    */
-  onTestFailed: (fn: OnTestFailedHandler) => void
+  onTestFailed: (fn: OnTestFailedHandler, timeout?: number) => void
 
   /**
    * Extract hooks on test failed
    */
-  onTestFinished: (fn: OnTestFinishedHandler) => void
+  onTestFinished: (fn: OnTestFinishedHandler, timeout?: number) => void
 
   /**
    * Mark tests as skipped. All execution after this call will be skipped.
@@ -625,8 +625,8 @@ export interface TaskContext<Task extends Test = Test> {
 export type ExtendedContext<T extends Test> = TaskContext<T> &
   TestContext
 
-export type OnTestFailedHandler = (result: TaskResult) => Awaitable<void>
-export type OnTestFinishedHandler = (result: TaskResult) => Awaitable<void>
+export type OnTestFailedHandler = (context: ExtendedContext<Test>) => Awaitable<void>
+export type OnTestFinishedHandler = (context: ExtendedContext<Test>) => Awaitable<void>
 
 export interface TaskHook<HookListener> {
   (fn: HookListener, timeout?: number): void