feat: implement aroundEach and aroundAll hooks

Author: sheremet-va

packages/runner/src/hooks.ts

@@ -1,6 +1,7 @@
 import type {
   AfterAllListener,
   AfterEachListener,
+  AroundEachListener,
   BeforeAllListener,
   BeforeEachListener,
   OnTestFailedHandler,
@@ -266,6 +267,46 @@ export const onTestFinished: TaskHook<OnTestFinishedHandler> = createTestHook(
   },
 )
 
+/**
+ * Registers a callback function that wraps around each test within the current suite.
+ * The callback receives a `runTest` function that must be called to run the test.
+ * This hook is useful for scenarios where you need to wrap tests in a context (e.g., database transactions).
+ *
+ * **Note:** When multiple `aroundEach` hooks are registered, they are nested inside each other.
+ * The first registered hook is the outermost wrapper.
+ *
+ * @param {Function} fn - The callback function that wraps the test. Must call `runTest()` to run the test.
+ * @param {number} [timeout] - Optional timeout in milliseconds for the hook. If not provided, the default hook timeout from the runner's configuration is used.
+ * @returns {void}
+ * @example
+ * ```ts
+ * // Example of using aroundEach to wrap tests in a database transaction
+ * aroundEach(async (runTest) => {
+ *   await database.beginTransaction();
+ *   await runTest(); // Run the test
+ *   await database.rollback();
+ * });
+ * ```
+ */
+export function aroundEach<ExtraContext = object>(
+  fn: AroundEachListener<ExtraContext>,
+  timeout?: number,
+): void {
+  assertTypes(fn, '"aroundEach" callback', ['function'])
+  const runner = getRunner()
+  return getCurrentSuite<ExtraContext>().on(
+    'aroundEach',
+    withTimeout(
+      fn,
+      // TODO: what should be the timeout? it runs _every_ hook inside (+a test)
+      timeout ?? (getDefaultHookTimeout() + runner.config.testTimeout + getDefaultHookTimeout()),
+      true,
+      new Error('STACK_TRACE_ERROR'),
+      ([, context], error) => abortContextSignal(context, error),
+    ),
+  )
+}
+
 function createTestHook<T>(
   name: string,
   handler: (test: TaskPopulated, handler: T, timeout?: number) => void,

packages/runner/src/index.ts

@@ -2,6 +2,7 @@ export { recordArtifact } from './artifact'
 export {
   afterAll,
   afterEach,
+  aroundEach,
   beforeAll,
   beforeEach,
   onTestFailed,

packages/runner/src/run.ts

@@ -2,6 +2,7 @@ import type { Awaitable, TestError } from '@vitest/utils'
 import type { DiffOptions } from '@vitest/utils/diff'
 import type { FileSpecification, VitestRunner } from './types/runner'
 import type {
+  AroundEachListener,
   File,
   SequenceHooks,
   Suite,
@@ -217,6 +218,51 @@ export async function callSuiteHook<T extends keyof SuiteHooks>(
   return callbacks
 }
 
+function getAroundEachHooks(suite: Suite): AroundEachListener[] {
+  const hooks: AroundEachListener[] = []
+  const parentSuite: Suite | null = 'filepath' in suite ? null : suite.suite || suite.file
+  if (parentSuite) {
+    hooks.push(...getAroundEachHooks(parentSuite))
+  }
+  hooks.push(...getHooks(suite).aroundEach)
+  return hooks
+}
+
+async function callAroundEachHooks(
+  suite: Suite,
+  test: Test,
+  runTest: () => Promise<void>,
+): Promise<void> {
+  const hooks = getAroundEachHooks(suite)
+
+  if (!hooks.length) {
+    await runTest()
+    return
+  }
+
+  const runNextHook = async (index: number): Promise<void> => {
+    if (index >= hooks.length) {
+      return runTest()
+    }
+
+    const hook = hooks[index]
+    let useCalled = false
+    const use = () => {
+      useCalled = true
+      return runNextHook(index + 1)
+    }
+    await hook(use, test.context, suite)
+    if (!useCalled) {
+      throw new Error(
+        'The `runTest()` callback was not called in the `aroundEach` hook. '
+        + 'Make sure to call `runTest()` to run the test.',
+      )
+    }
+  }
+
+  await runNextHook(0)
+}
+
 const packs = new Map<string, [TaskResult | undefined, TaskMeta]>()
 const eventsPacks: [string, TaskUpdateEvent, undefined][] = []
 const pendingTasksUpdates: Promise<void>[] = []
@@ -365,93 +411,95 @@ export async function runTest(test: Test, runner: VitestRunner): Promise<void> {
     const retry = getRetryCount(test.retry)
     for (let retryCount = 0; retryCount <= retry; retryCount++) {
       let beforeEachCleanups: unknown[] = []
-      try {
-        await runner.onBeforeTryTask?.(test, {
-          retry: retryCount,
-          repeats: repeatCount,
-        })
-
-        test.result.repeatCount = repeatCount
+      await callAroundEachHooks(suite, test, async () => {
+        try {
+          await runner.onBeforeTryTask?.(test, {
+            retry: retryCount,
+            repeats: repeatCount,
+          })
+
+          test.result!.repeatCount = repeatCount
+
+          beforeEachCleanups = await $('test.beforeEach', () => callSuiteHook(
+            suite,
+            test,
+            'beforeEach',
+            runner,
+            [test.context, suite],
+          ))
+
+          if (runner.runTask) {
+            await $('test.callback', () => runner.runTask!(test))
+          }
+          else {
+            const fn = getFn(test)
+            if (!fn) {
+              throw new Error(
+                'Test function is not found. Did you add it using `setFn`?',
+              )
+            }
+            await $('test.callback', () => fn())
+          }
 
-        beforeEachCleanups = await $('test.beforeEach', () => callSuiteHook(
-          suite,
-          test,
-          'beforeEach',
-          runner,
-          [test.context, suite],
-        ))
+          await runner.onAfterTryTask?.(test, {
+            retry: retryCount,
+            repeats: repeatCount,
+          })
 
-        if (runner.runTask) {
-          await $('test.callback', () => runner.runTask!(test))
-        }
-        else {
-          const fn = getFn(test)
-          if (!fn) {
-            throw new Error(
-              'Test function is not found. Did you add it using `setFn`?',
-            )
+          if (test.result!.state !== 'fail') {
+            if (!test.repeats) {
+              test.result!.state = 'pass'
+            }
+            else if (test.repeats && retry === retryCount) {
+              test.result!.state = 'pass'
+            }
           }
-          await $('test.callback', () => fn())
+        }
+        catch (e) {
+          failTask(test.result!, e, runner.config.diffOptions)
         }
 
-        await runner.onAfterTryTask?.(test, {
-          retry: retryCount,
-          repeats: repeatCount,
-        })
+        try {
+          await runner.onTaskFinished?.(test)
+        }
+        catch (e) {
+          failTask(test.result!, e, runner.config.diffOptions)
+        }
 
-        if (test.result.state !== 'fail') {
-          if (!test.repeats) {
-            test.result.state = 'pass'
-          }
-          else if (test.repeats && retry === retryCount) {
-            test.result.state = 'pass'
+        try {
+          await $('test.afterEach', () => callSuiteHook(suite, test, 'afterEach', runner, [
+            test.context,
+            suite,
+          ]))
+          if (beforeEachCleanups.length) {
+            await $('test.cleanup', () => callCleanupHooks(runner, beforeEachCleanups))
           }
+          await callFixtureCleanup(test.context)
         }
-      }
-      catch (e) {
-        failTask(test.result, e, runner.config.diffOptions)
-      }
-
-      try {
-        await runner.onTaskFinished?.(test)
-      }
-      catch (e) {
-        failTask(test.result, e, runner.config.diffOptions)
-      }
-
-      try {
-        await $('test.afterEach', () => callSuiteHook(suite, test, 'afterEach', runner, [
-          test.context,
-          suite,
-        ]))
-        if (beforeEachCleanups.length) {
-          await $('test.cleanup', () => callCleanupHooks(runner, beforeEachCleanups))
+        catch (e) {
+          failTask(test.result!, e, runner.config.diffOptions)
         }
-        await callFixtureCleanup(test.context)
-      }
-      catch (e) {
-        failTask(test.result, e, runner.config.diffOptions)
-      }
 
-      if (test.onFinished?.length) {
-        await $('test.onFinished', () => callTestHooks(runner, test, test.onFinished!, 'stack'))
-      }
+        if (test.onFinished?.length) {
+          await $('test.onFinished', () => callTestHooks(runner, test, test.onFinished!, 'stack'))
+        }
 
-      if (test.result.state === 'fail' && test.onFailed?.length) {
-        await $('test.onFailed', () => callTestHooks(
-          runner,
-          test,
-          test.onFailed!,
-          runner.config.sequence.hooks,
-        ))
-      }
+        if (test.result!.state === 'fail' && test.onFailed?.length) {
+          await $('test.onFailed', () => callTestHooks(
+            runner,
+            test,
+            test.onFailed!,
+            runner.config.sequence.hooks,
+          ))
+        }
 
-      test.onFailed = undefined
-      test.onFinished = undefined
+        test.onFailed = undefined
+        test.onFinished = undefined
 
-      await runner.onAfterRetryTask?.(test, {
-        retry: retryCount,
-        repeats: repeatCount,
+        await runner.onAfterRetryTask?.(test, {
+          retry: retryCount,
+          repeats: repeatCount,
+        })
       })
 
       // skipped with new PendingError

packages/runner/src/suite.ts

@@ -247,6 +247,7 @@ export function createSuiteHooks(): SuiteHooks {
     afterAll: [],
     beforeEach: [],
     afterEach: [],
+    aroundEach: [],
   }
 }
 

packages/runner/src/types.ts

@@ -9,6 +9,7 @@ export type {
 export type {
   AfterAllListener,
   AfterEachListener,
+  AroundEachListener,
   BeforeAllListener,
   BeforeEachListener,
   File,

packages/runner/src/types/tasks.ts

@@ -688,11 +688,20 @@ export interface AfterEachListener<ExtraContext = object> {
   ): Awaitable<unknown>
 }
 
+export interface AroundEachListener<ExtraContext = object> {
+  (
+    runTest: () => Promise<void>,
+    context: TestContext & ExtraContext,
+    suite: Readonly<Suite>
+  ): Awaitable<unknown>
+}
+
 export interface SuiteHooks<ExtraContext = object> {
   beforeAll: BeforeAllListener[]
   afterAll: AfterAllListener[]
   beforeEach: BeforeEachListener<ExtraContext>[]
   afterEach: AfterEachListener<ExtraContext>[]
+  aroundEach: AroundEachListener<ExtraContext>[]
 }
 
 export interface TaskCustomOptions extends TestOptions {

packages/vitest/src/public/index.ts

@@ -97,6 +97,7 @@ export type {
 export {
   afterAll,
   afterEach,
+  aroundEach,
   beforeAll,
   beforeEach,
   describe,