feat(vitest): add onTestFinished hook (#5128)

Author: sheremet-va

docs/api/index.md

@@ -853,6 +853,10 @@ afterEach(async () => {
 
 Here, the `afterEach` ensures that testing data is cleared after each test runs.
 
+::: tip
+Vitest 1.3.0 added [`onTestFinished`](##ontestfinished-1-3-0) hook. You can call it during the test execution to cleanup any state after the test has finished running.
+:::
+
 ### beforeAll
 
 - **Type:** `beforeAll(fn: () => Awaitable<void>, timeout?: number)`
@@ -906,3 +910,96 @@ afterAll(async () => {
 ```
 
 Here the `afterAll` ensures that `stopMocking` method is called after all tests run.
+
+## Test Hooks
+
+Vitest provides a few hooks that you can call _during_ the test execution to cleanup the state when the test has finished runnning.
+
+::: warning
+These hooks will throw an error if they are called outside of the test body.
+:::
+
+### onTestFinished <Badge type="info">1.3.0+</Badge>
+
+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.
+
+```ts
+import { onTestFinished, test } from 'vitest'
+
+test('performs a query', () => {
+  const db = connectDb()
+  onTestFinished(() => db.close())
+  db.query('SELECT * FROM users')
+})
+```
+
+::: warning
+If you are running tests concurrently, you should always use `onTestFinished` hook from the test context since Vitest doesn't track concurrent tests in global hooks:
+
+```ts
+import { test } from 'vitest'
+
+test.concurrent('performs a query', (t) => {
+  const db = connectDb()
+  t.onTestFinished(() => db.close())
+  db.query('SELECT * FROM users')
+})
+```
+:::
+
+This hook is particularly useful when creating reusable logic:
+
+```ts
+// this can be in a separate file
+function getTestDb() {
+  const db = connectMockedDb()
+  onTestFinished(() => db.close())
+  return db
+}
+
+test('performs a user query', async () => {
+  const db = getTestDb()
+  expect(
+    await db.query('SELECT * from users').perform()
+  ).toEqual([])
+})
+
+test('performs an organization query', async () => {
+  const db = getTestDb()
+  expect(
+    await db.query('SELECT * from organizations').perform()
+  ).toEqual([])
+})
+```
+
+### 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.
+
+```ts
+import { onTestFailed, test } from 'vitest'
+
+test('performs a query', () => {
+  const db = connectDb()
+  onTestFailed((e) => {
+    console.log(e.result.errors)
+  })
+  db.query('SELECT * FROM users')
+})
+```
+
+::: warning
+If you are running tests concurrently, you should always use `onTestFailed` hook from the test context since Vitest doesn't track concurrent tests in global hooks:
+
+```ts
+import { test } from 'vitest'
+
+test.concurrent('performs a query', (t) => {
+  const db = connectDb()
+  onTestFailed((result) => {
+    console.log(result.errors)
+  })
+  db.query('SELECT * FROM users')
+})
+```
+:::

packages/runner/src/context.ts

@@ -59,6 +59,11 @@ export function createTestContext<T extends Test | Custom>(test: T, runner: Vite
     test.onFailed.push(fn)
   }
 
+  context.onTestFinished = (fn) => {
+    test.onFinished ||= []
+    test.onFinished.push(fn)
+  }
+
   return runner.extendTaskContext?.(context) as ExtendedContext<T> || context
 }
 

packages/runner/src/hooks.ts

@@ -1,4 +1,4 @@
-import type { OnTestFailedHandler, SuiteHooks, TaskPopulated } from './types'
+import type { OnTestFailedHandler, OnTestFinishedHandler, SuiteHooks, TaskPopulated } from './types'
 import { getCurrentSuite, getRunner } from './suite'
 import { getCurrentTest } from './test-state'
 import { withTimeout } from './context'
@@ -27,13 +27,18 @@ export const onTestFailed = createTestHook<OnTestFailedHandler>('onTestFailed',
   test.onFailed.push(handler)
 })
 
+export const onTestFinished = createTestHook<OnTestFinishedHandler>('onTestFinished', (test, handler) => {
+  test.onFinished ||= []
+  test.onFinished.push(handler)
+})
+
 function createTestHook<T>(name: string, handler: (test: TaskPopulated, handler: T) => void) {
   return (fn: T) => {
     const current = getCurrentTest()
 
     if (!current)
       throw new Error(`Hook ${name}() can only be called inside a test`)
 
-    handler(current, fn)
+    return handler(current, fn)
   }
 }

packages/runner/src/index.ts

@@ -1,6 +1,6 @@
 export { startTests, updateTask } from './run'
 export { test, it, describe, suite, getCurrentSuite, createTaskCollector } from './suite'
-export { beforeAll, beforeEach, afterAll, afterEach, onTestFailed } from './hooks'
+export { beforeAll, beforeEach, afterAll, afterEach, onTestFailed, onTestFinished } from './hooks'
 export { setFn, getFn, getHooks, setHooks } from './map'
 export { getCurrentTest } from './test-state'
 export { processError } from '@vitest/utils/error'

packages/runner/src/run.ts

@@ -210,8 +210,21 @@ export async function runTest(test: Test | Custom, runner: VitestRunner) {
     }
   }
 
-  if (test.result.state === 'fail')
-    await Promise.all(test.onFailed?.map(fn => fn(test.result!)) || [])
+  try {
+    await Promise.all(test.onFinished?.map(fn => fn(test.result!)) || [])
+  }
+  catch (e) {
+    failTask(test.result, e, runner.config.diffOptions)
+  }
+
+  if (test.result.state === 'fail') {
+    try {
+      await Promise.all(test.onFailed?.map(fn => fn(test.result!)) || [])
+    }
+    catch (e) {
+      failTask(test.result, e, runner.config.diffOptions)
+    }
+  }
 
   // if test is marked to be failed, flip the result
   if (test.fails) {

packages/runner/src/types/tasks.ts

@@ -26,6 +26,7 @@ export interface TaskPopulated extends TaskBase {
   result?: TaskResult
   fails?: boolean
   onFailed?: OnTestFailedHandler[]
+  onFinished?: OnTestFinishedHandler[]
   /**
    * Store promises (from async expects) to wait for them before finishing the test
    */
@@ -296,6 +297,11 @@ export interface TaskContext<Task extends Custom | Test = Custom | Test> {
    */
   onTestFailed: (fn: OnTestFailedHandler) => void
 
+  /**
+   * Extract hooks on test failed
+   */
+  onTestFinished: (fn: OnTestFinishedHandler) => void
+
   /**
    * Mark tests as skipped. All execution after this call will be skipped.
    */
@@ -305,6 +311,7 @@ export interface TaskContext<Task extends Custom | Test = Custom | Test> {
 export type ExtendedContext<T extends Custom | Test> = TaskContext<T> & TestContext
 
 export type OnTestFailedHandler = (result: TaskResult) => Awaitable<void>
+export type OnTestFinishedHandler = (result: TaskResult) => Awaitable<void>
 
 export type SequenceHooks = 'stack' | 'list' | 'parallel'
 export type SequenceSetupFiles = 'list' | 'parallel'

packages/vitest/src/index.ts

@@ -8,6 +8,7 @@ export {
   afterAll,
   afterEach,
   onTestFailed,
+  onTestFinished,
 } from '@vitest/runner'
 export { bench } from './runtime/benchmark'
 

test/core/test/on-finished.test.ts

@@ -0,0 +1,43 @@
+import { expect, it, onTestFinished } from 'vitest'
+
+const collected: any[] = []
+
+it('on-finished regular', () => {
+  collected.push(1)
+  onTestFinished(() => {
+    collected.push(3)
+  })
+  collected.push(2)
+})
+
+it('on-finished context', (t) => {
+  collected.push(4)
+  t.onTestFinished(() => {
+    collected.push(6)
+  })
+  collected.push(5)
+})
+
+it.fails('failed finish', () => {
+  collected.push(7)
+  onTestFinished(() => {
+    collected.push(9)
+  })
+  collected.push(8)
+  expect.fail('failed')
+  collected.push(null)
+})
+
+it.fails('failed finish context', (t) => {
+  collected.push(10)
+  t.onTestFinished(() => {
+    collected.push(12)
+  })
+  collected.push(11)
+  expect.fail('failed')
+  collected.push(null)
+})
+
+it('after', () => {
+  expect(collected).toEqual([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12])
+})