feat(vitest): add onTestsRerun method to global setup context (#6803)

Author: sheremet-va

docs/config/index.md

@@ -1089,6 +1089,18 @@ inject('wsPort') === 3000
 ```
 :::
 
+Since Vitest 2.2.0, you can define a custom callback function to be called when Vitest reruns tests. If the function is asynchronous, the runner will wait for it to complete before executing the tests.
+
+```ts
+import type { GlobalSetupContext } from 'vitest/node'
+
+export default function setup({ onTestsRerun }: GlobalSetupContext) {
+  onTestsRerun(async () => {
+    await restartDb()
+  })
+}
+```
+
 ### forceRerunTriggers<NonProjectOption />
 
 - **Type**: `string[]`

packages/vitest/src/node/core.ts

@@ -3,7 +3,7 @@ import type { Writable } from 'node:stream'
 import type { ViteDevServer } from 'vite'
 import type { defineWorkspace } from 'vitest/config'
 import type { SerializedCoverageConfig } from '../runtime/config'
-import type { ArgumentsType, OnServerRestartHandler, ProvidedContext, UserConsoleLog } from '../types/general'
+import type { ArgumentsType, OnServerRestartHandler, OnTestsRerunHandler, ProvidedContext, UserConsoleLog } from '../types/general'
 import type { ProcessPool, WorkspaceSpec } from './pool'
 import type { TestSpecification } from './spec'
 import type { ResolvedConfig, UserConfig, VitestRunMode } from './types/config'
@@ -104,6 +104,7 @@ export class Vitest {
   private _onClose: (() => Awaited<unknown>)[] = []
   private _onSetServer: OnServerRestartHandler[] = []
   private _onCancelListeners: ((reason: CancelReason) => Promise<void> | void)[] = []
+  private _onUserTestsRerun: OnTestsRerunHandler[] = []
 
   async setServer(options: UserConfig, server: ViteDevServer, cliOptions: UserConfig) {
     this.unregisterWatcher?.()
@@ -119,6 +120,7 @@ export class Vitest {
     this.coverageProvider = undefined
     this.runningPromise = undefined
     this._cachedSpecs.clear()
+    this._onUserTestsRerun = []
 
     const resolved = resolveConfig(this.mode, options, server.config, this.logger)
 
@@ -695,7 +697,10 @@ export class Vitest {
       files = files.filter(file => filteredFiles.some(f => f[1] === file))
     }
 
-    await this.report('onWatcherRerun', files, trigger)
+    await Promise.all([
+      this.report('onWatcherRerun', files, trigger),
+      ...this._onUserTestsRerun.map(fn => fn(files)),
+    ])
     await this.runFiles(files.flatMap(file => this.getProjectsByTestFile(file)), allTestsRun)
 
     await this.report('onWatcherStart', this.state.getFiles(files))
@@ -817,7 +822,10 @@ export class Vitest {
 
       const triggerIds = new Set(triggerId.map(id => relative(this.config.root, id)))
       const triggerLabel = Array.from(triggerIds).join(', ')
-      await this.report('onWatcherRerun', files, triggerLabel)
+      await Promise.all([
+        this.report('onWatcherRerun', files, triggerLabel),
+        ...this._onUserTestsRerun.map(fn => fn(files)),
+      ])
 
       await this.runFiles(files.flatMap(file => this.getProjectsByTestFile(file)), false)
 
@@ -1154,4 +1162,8 @@ export class Vitest {
   onClose(fn: () => void) {
     this._onClose.push(fn)
   }
+
+  onTestsRerun(fn: OnTestsRerunHandler): void {
+    this._onUserTestsRerun.push(fn)
+  }
 }

packages/vitest/src/node/globalSetup.ts

@@ -1,5 +1,5 @@
 import type { ViteNodeRunner } from 'vite-node/client'
-import type { ProvidedContext } from '../types/general'
+import type { OnTestsRerunHandler, ProvidedContext } from '../types/general'
 import type { ResolvedConfig } from './types/config'
 import { toArray } from '@vitest/utils'
 
@@ -9,6 +9,7 @@ export interface GlobalSetupContext {
     key: T,
     value: ProvidedContext[T]
   ) => void
+  onTestsRerun: (cb: OnTestsRerunHandler) => void
 }
 
 export interface GlobalSetupFile {

packages/vitest/src/node/workspace.ts

@@ -170,6 +170,7 @@ export class WorkspaceProject {
       const teardown = await globalSetupFile.setup?.({
         provide: (key, value) => this.provide(key, value),
         config: this.config,
+        onTestsRerun: cb => this.ctx.onTestsRerun(cb),
       })
       if (teardown == null || !!globalSetupFile.teardown) {
         continue

packages/vitest/src/public/node.ts

@@ -126,7 +126,10 @@ export type {
   TscErrorInfo as TypeCheckErrorInfo,
 } from '../typecheck/types'
 
-export type { OnServerRestartHandler } from '../types/general'
+export type {
+  OnServerRestartHandler,
+  OnTestsRerunHandler,
+} from '../types/general'
 
 export { createDebugger } from '../utils/debugger'
 

packages/vitest/src/types/general.ts

@@ -46,5 +46,5 @@ export interface ModuleGraphData {
 }
 
 export type OnServerRestartHandler = (reason?: string) => Promise<void> | void
-
+export type OnTestsRerunHandler = (testFiles: string[]) => Promise<void> | void
 export interface ProvidedContext {}

test/watch/fixtures/global-setup.ts

@@ -0,0 +1,15 @@
+import { GlobalSetupContext } from 'vitest/node';
+
+const calls: string[] = [];
+
+(globalThis as any).__CALLS = calls
+
+export default ({ onTestsRerun }: GlobalSetupContext) => {
+  calls.push('start')
+  onTestsRerun(() => {
+    calls.push('rerun')
+  })
+  return () => {
+    calls.push('end')
+  }
+}

test/watch/fixtures/vitest.config.ts

@@ -18,5 +18,9 @@ export default defineConfig({
     forceRerunTriggers: [
       '**/force-watch/**',
     ],
+
+    globalSetup: process.env.TEST_GLOBAL_SETUP
+      ? './global-setup.ts'
+      : undefined,
   },
 })

test/watch/test/global-setup-rerun.test.ts

@@ -0,0 +1,50 @@
+import { expect, test } from 'vitest'
+import { editFile, runVitest } from '../../test-utils'
+
+const testFile = 'fixtures/math.test.ts'
+
+test('global setup calls hooks correctly when file changes', async () => {
+  process.env.TEST_GLOBAL_SETUP = 'true'
+  const { vitest, ctx } = await runVitest({
+    root: 'fixtures',
+    watch: true,
+    include: ['math.test.ts'],
+  })
+
+  await vitest.waitForStdout('Waiting for file changes')
+
+  const calls = (globalThis as any).__CALLS as string[]
+  expect(calls).toEqual(['start'])
+
+  editFile(testFile, testFileContent => `${testFileContent}\n\n`)
+
+  await vitest.waitForStdout('RERUN')
+  expect(calls).toEqual(['start', 'rerun'])
+
+  await ctx?.close()
+
+  expect(calls).toEqual(['start', 'rerun', 'end'])
+})
+
+test('global setup calls hooks correctly with a manual rerun', async () => {
+  process.env.TEST_GLOBAL_SETUP = 'true'
+  const { vitest, ctx } = await runVitest({
+    root: 'fixtures',
+    watch: true,
+    include: ['math.test.ts'],
+  })
+
+  await vitest.waitForStdout('Waiting for file changes')
+
+  const calls = (globalThis as any).__CALLS as string[]
+  expect(calls).toEqual(['start'])
+
+  vitest.write('r')
+
+  await vitest.waitForStdout('RERUN')
+  expect(calls).toEqual(['start', 'rerun'])
+
+  await ctx?.close()
+
+  expect(calls).toEqual(['start', 'rerun', 'end'])
+})