feat: provide ctx.signal (#7878)

Author: sheremet-va

docs/guide/test-context.md

@@ -79,6 +79,21 @@ it('math is hard', ({ skip, mind }) => {
 })
 ```
 
+#### `context.signal` <Version>3.2.0</Version> {#context-signal}
+
+An [`AbortSignal`](https://developer.mozilla.org/en-US/docs/Web/API/AbortSignal) that can be aborted by Vitest. The signal is aborted in these situations:
+
+- Test times out
+- User manually cancelled the test run with Ctrl+C
+- [`vitest.cancelCurrentRun`](/advanced/api/vitest#cancelcurrentrun) was called programmatically
+- Another test failed in parallel and the [`bail`](/config/#bail) flag is set
+
+```ts
+it('stop request when test times out', async ({ signal }) => {
+  await fetch('/resource', { signal })
+}, 2000)
+```
+
 #### `onTestFailed`
 
 The [`onTestFailed`](/api/#ontestfailed) hook bound to the current test. This API is useful if you are running tests concurrently and need to have a special handling only for this specific test.

packages/browser/src/client/tester/runner.ts

@@ -62,8 +62,8 @@ export function createBrowserRunner(
         const currentFailures = 1 + previousFailures
 
         if (currentFailures >= this.config.bail) {
-          rpc().onCancel('test-failure')
-          this.onCancel('test-failure')
+          rpc().cancelCurrentRun('test-failure')
+          this.cancel('test-failure')
         }
       }
     }
@@ -81,8 +81,8 @@ export function createBrowserRunner(
       }
     }
 
-    onCancel = (reason: CancelReason) => {
-      super.onCancel?.(reason)
+    cancel = (reason: CancelReason) => {
+      super.cancel?.(reason)
       globalChannel.postMessage({ type: 'cancel', reason })
     }
 
@@ -196,7 +196,7 @@ export async function initiateRunner(
   cachedRunner = runner
 
   onCancel.then((reason) => {
-    runner.onCancel?.(reason)
+    runner.cancel?.(reason)
   })
 
   const [diffOptions] = await Promise.all([

packages/browser/src/node/rpc.ts

@@ -202,7 +202,7 @@ export function setupBrowserRpc(globalServer: ParentBrowserProject, defaultMocke
           const mod = globalServer.vite.moduleGraph.getModuleById(id)
           return mod?.transformResult?.map
         },
-        onCancel(reason) {
+        cancelCurrentRun(reason) {
           vitest.cancelCurrentRun(reason)
         },
         async resolveId(id, importer) {

packages/browser/src/node/types.ts

@@ -21,7 +21,7 @@ export interface WebSocketBrowserHandlers {
   onCollected: (method: TestExecutionMethod, files: RunnerTestFile[]) => Promise<void>
   onTaskUpdate: (method: TestExecutionMethod, packs: TaskResultPack[], events: TaskEventPack[]) => void
   onAfterSuiteRun: (meta: AfterSuiteRunMeta) => void
-  onCancel: (reason: CancelReason) => void
+  cancelCurrentRun: (reason: CancelReason) => void
   getCountOfFailedTests: () => number
   readSnapshotFile: (id: string) => Promise<string | null>
   saveSnapshotFile: (id: string, content: string) => Promise<void>

packages/runner/src/context.ts

@@ -5,6 +5,7 @@ import type {
   SuiteCollector,
   Test,
   TestContext,
+  WriteableTestContext,
 } from './types/tasks'
 import { getSafeTimers } from '@vitest/utils'
 import { PendingError } from './errors'
@@ -36,6 +37,7 @@ export function withTimeout<T extends (...args: any[]) => any>(
   timeout: number,
   isHook = false,
   stackTraceError?: Error,
+  onTimeout?: (args: T extends (...args: infer A) => any ? A : never, error: Error) => void,
 ): T {
   if (timeout <= 0 || timeout === Number.POSITIVE_INFINITY) {
     return fn
@@ -58,7 +60,9 @@ export function withTimeout<T extends (...args: any[]) => any>(
       timer.unref?.()
 
       function rejectTimeoutError() {
-        reject_(makeTimeoutError(isHook, timeout, stackTraceError))
+        const error = makeTimeoutError(isHook, timeout, stackTraceError)
+        onTimeout?.(args, error)
+        reject_(error)
       }
 
       function resolve(result: unknown) {
@@ -102,14 +106,35 @@ export function withTimeout<T extends (...args: any[]) => any>(
   }) as T
 }
 
+const abortControllers = new WeakMap<TestContext, AbortController>()
+
+export function abortIfTimeout([context]: [TestContext?], error: Error): void {
+  if (context) {
+    abortContextSignal(context, error)
+  }
+}
+
+export function abortContextSignal(context: TestContext, error: Error): void {
+  const abortController = abortControllers.get(context)
+  abortController?.abort(error)
+}
+
 export function createTestContext(
   test: Test,
   runner: VitestRunner,
 ): TestContext {
   const context = function () {
     throw new Error('done() callback is deprecated, use promise instead')
-  } as unknown as TestContext
+  } as unknown as WriteableTestContext
+
+  let abortController = abortControllers.get(context)
+
+  if (!abortController) {
+    abortController = new AbortController()
+    abortControllers.set(context, abortController)
+  }
 
+  context.signal = abortController.signal
   context.task = test
 
   context.skip = (condition?: boolean | string, note?: string): never => {
@@ -129,14 +154,26 @@ export function createTestContext(
   context.onTestFailed = (handler, timeout) => {
     test.onFailed ||= []
     test.onFailed.push(
-      withTimeout(handler, timeout ?? runner.config.hookTimeout, true, new Error('STACK_TRACE_ERROR')),
+      withTimeout(
+        handler,
+        timeout ?? runner.config.hookTimeout,
+        true,
+        new Error('STACK_TRACE_ERROR'),
+        (_, error) => abortController.abort(error),
+      ),
     )
   }
 
   context.onTestFinished = (handler, timeout) => {
     test.onFinished ||= []
     test.onFinished.push(
-      withTimeout(handler, timeout ?? runner.config.hookTimeout, true, new Error('STACK_TRACE_ERROR')),
+      withTimeout(
+        handler,
+        timeout ?? runner.config.hookTimeout,
+        true,
+        new Error('STACK_TRACE_ERROR'),
+        (_, error) => abortController.abort(error),
+      ),
     )
   }
 

packages/runner/src/errors.ts

@@ -1,3 +1,4 @@
+import type { CancelReason } from './types/runner'
 import type { TaskBase } from './types/tasks'
 
 export class PendingError extends Error {
@@ -9,3 +10,12 @@ export class PendingError extends Error {
     this.taskId = task.id
   }
 }
+
+export class TestRunAbortError extends Error {
+  public name = 'TestRunAbortError'
+  public reason: CancelReason
+  constructor(message: string, reason: CancelReason) {
+    super(message)
+    this.reason = reason
+  }
+}

packages/runner/src/hooks.ts

@@ -7,9 +7,10 @@ import type {
   OnTestFinishedHandler,
   TaskHook,
   TaskPopulated,
+  TestContext,
 } from './types/tasks'
 import { assertTypes } from '@vitest/utils'
-import { withTimeout } from './context'
+import { abortContextSignal, abortIfTimeout, withTimeout } from './context'
 import { withFixtures } from './fixture'
 import { getCurrentSuite, getRunner } from './suite'
 import { getCurrentTest } from './test-state'
@@ -21,7 +22,7 @@ function getDefaultHookTimeout() {
 const CLEANUP_TIMEOUT_KEY = Symbol.for('VITEST_CLEANUP_TIMEOUT')
 const CLEANUP_STACK_TRACE_KEY = Symbol.for('VITEST_CLEANUP_STACK_TRACE')
 
-export function getBeforeHookCleanupCallback(hook: Function, result: any): Function | undefined {
+export function getBeforeHookCleanupCallback(hook: Function, result: any, context?: TestContext): Function | undefined {
   if (typeof result === 'function') {
     const timeout
       = CLEANUP_TIMEOUT_KEY in hook && typeof hook[CLEANUP_TIMEOUT_KEY] === 'number'
@@ -31,7 +32,17 @@ export function getBeforeHookCleanupCallback(hook: Function, result: any): Funct
       = CLEANUP_STACK_TRACE_KEY in hook && hook[CLEANUP_STACK_TRACE_KEY] instanceof Error
         ? hook[CLEANUP_STACK_TRACE_KEY]
         : undefined
-    return withTimeout(result, timeout, true, stackTraceError)
+    return withTimeout(
+      result,
+      timeout,
+      true,
+      stackTraceError,
+      (_, error) => {
+        if (context) {
+          abortContextSignal(context, error)
+        }
+      },
+    )
   }
 }
 
@@ -136,6 +147,7 @@ export function beforeEach<ExtraContext = object>(
         timeout ?? getDefaultHookTimeout(),
         true,
         stackTraceError,
+        abortIfTimeout,
       ),
       {
         [CLEANUP_TIMEOUT_KEY]: timeout,
@@ -174,6 +186,7 @@ export function afterEach<ExtraContext = object>(
       timeout ?? getDefaultHookTimeout(),
       true,
       new Error('STACK_TRACE_ERROR'),
+      abortIfTimeout,
     ),
   )
 }
@@ -206,6 +219,7 @@ export const onTestFailed: TaskHook<OnTestFailedHandler> = createTestHook(
         timeout ?? getDefaultHookTimeout(),
         true,
         new Error('STACK_TRACE_ERROR'),
+        abortIfTimeout,
       ),
     )
   },
@@ -244,6 +258,7 @@ export const onTestFinished: TaskHook<OnTestFinishedHandler> = createTestHook(
         timeout ?? getDefaultHookTimeout(),
         true,
         new Error('STACK_TRACE_ERROR'),
+        abortIfTimeout,
       ),
     )
   },

packages/runner/src/run.ts

@@ -15,15 +15,17 @@ import type {
   TaskUpdateEvent,
   Test,
   TestContext,
+  WriteableTestContext,
 } from './types/tasks'
 import { shuffle } from '@vitest/utils'
 import { processError } from '@vitest/utils/error'
 import { collectTests } from './collect'
-import { PendingError } from './errors'
+import { abortContextSignal } from './context'
+import { PendingError, TestRunAbortError } from './errors'
 import { callFixtureCleanup } from './fixture'
 import { getBeforeHookCleanupCallback } from './hooks'
 import { getFn, getHooks } from './map'
-import { setCurrentTest } from './test-state'
+import { addRunningTest, getRunningTests, setCurrentTest } from './test-state'
 import { limitConcurrency } from './utils/limit-concurrency'
 import { partitionSuiteChildren } from './utils/suite'
 import { hasFailed, hasTests } from './utils/tasks'
@@ -87,12 +89,14 @@ async function callTestHooks(
     return
   }
 
+  const context = test.context as WriteableTestContext
+
   const onTestFailed = test.context.onTestFailed
   const onTestFinished = test.context.onTestFinished
-  test.context.onTestFailed = () => {
+  context.onTestFailed = () => {
     throw new Error(`Cannot call "onTestFailed" inside a test hook.`)
   }
-  test.context.onTestFinished = () => {
+  context.onTestFinished = () => {
     throw new Error(`Cannot call "onTestFinished" inside a test hook.`)
   }
 
@@ -115,8 +119,8 @@ async function callTestHooks(
     }
   }
 
-  test.context.onTestFailed = onTestFailed
-  test.context.onTestFinished = onTestFinished
+  context.onTestFailed = onTestFailed
+  context.onTestFinished = onTestFinished
 }
 
 export async function callSuiteHook<T extends keyof SuiteHooks>(
@@ -145,7 +149,11 @@ export async function callSuiteHook<T extends keyof SuiteHooks>(
   }
 
   async function runHook(hook: Function) {
-    return getBeforeHookCleanupCallback(hook, await hook(...args))
+    return getBeforeHookCleanupCallback(
+      hook,
+      await hook(...args),
+      name === 'beforeEach' ? args[0] : undefined,
+    )
   }
 
   if (sequence === 'parallel') {
@@ -274,6 +282,7 @@ export async function runTest(test: Test, runner: VitestRunner): Promise<void> {
   }
   updateTask('test-prepare', test, runner)
 
+  const cleanupRunningTest = addRunningTest(test)
   setCurrentTest(test)
 
   const suite = test.suite || test.file
@@ -374,6 +383,7 @@ export async function runTest(test: Test, runner: VitestRunner): Promise<void> {
         }
         updateTask('test-finished', test, runner)
         setCurrentTest(undefined)
+        cleanupRunningTest()
         return
       }
 
@@ -405,6 +415,7 @@ export async function runTest(test: Test, runner: VitestRunner): Promise<void> {
     }
   }
 
+  cleanupRunningTest()
   setCurrentTest(undefined)
 
   test.result.duration = now() - start
@@ -588,21 +599,38 @@ export async function runFiles(files: File[], runner: VitestRunner): Promise<voi
 }
 
 export async function startTests(specs: string[] | FileSpecification[], runner: VitestRunner): Promise<File[]> {
-  const paths = specs.map(f => typeof f === 'string' ? f : f.filepath)
-  await runner.onBeforeCollect?.(paths)
+  const cancel = runner.cancel?.bind(runner)
+  // Ideally, we need to have an event listener for this, but only have a runner here.
+  // Adding another onCancel felt wrong (maybe it needs to be refactored)
+  runner.cancel = (reason) => {
+    // We intentionally create only one error since there is only one test run that can be cancelled
+    const error = new TestRunAbortError('The test run was aborted by the user.', reason)
+    getRunningTests().forEach(test =>
+      abortContextSignal(test.context, error),
+    )
+    return cancel?.(reason)
+  }
 
-  const files = await collectTests(specs, runner)
+  try {
+    const paths = specs.map(f => typeof f === 'string' ? f : f.filepath)
+    await runner.onBeforeCollect?.(paths)
 
-  await runner.onCollected?.(files)
-  await runner.onBeforeRunFiles?.(files)
+    const files = await collectTests(specs, runner)
 
-  await runFiles(files, runner)
+    await runner.onCollected?.(files)
+    await runner.onBeforeRunFiles?.(files)
 
-  await runner.onAfterRunFiles?.(files)
+    await runFiles(files, runner)
 
-  await finishSendTasksUpdate(runner)
+    await runner.onAfterRunFiles?.(files)
 
-  return files
+    await finishSendTasksUpdate(runner)
+
+    return files
+  }
+  finally {
+    runner.cancel = cancel
+  }
 }
 
 async function publicCollect(specs: string[] | FileSpecification[], runner: VitestRunner): Promise<File[]> {