feat(vitest): allow per-file and per-worker fixtures (#7704)

Co-authored-by: Ari Perkkiö <ari.perkkio@gmail.com>

Author: sheremet-va

docs/guide/test-context.md

@@ -373,6 +373,46 @@ describe('another type of schema', () => {
 })
 ```
 
+#### Per-Scope Context <Version>3.2.0</Version>
+
+You can define context that will be initiated once per file or a worker. It is initiated the same way as a regular fixture with an objects parameter:
+
+```ts
+import { test as baseTest } from 'vitest'
+
+export const test = baseTest.extend({
+  perFile: [
+    ({}, { use }) => use([]),
+    { scope: 'file' },
+  ],
+  perWorker: [
+    ({}, { use }) => use([]),
+    { scope: 'worker' },
+  ],
+})
+```
+
+The value is initialised the first time any test has accessed it, unless the fixture options have `auto: true` - in this case the value is initialised before any test has run.
+
+```ts
+const test = baseTest.extend({
+  perFile: [
+    ({}, { use }) => use([]),
+    {
+      scope: 'file',
+      // always run this hook before any test
+      auto: true
+    },
+  ],
+})
+```
+
+The `worker` scope will run the fixture once per worker. The number of running workers depends on various factors. By default, every file runs in a separate worker, so `file` and `worker` scopes work the same way.
+
+However, if you disable [isolation](/config/#isolate), then the number of workers is limited by the [`maxWorkers`](/config/#maxworkers) or [`poolOptions`](/config/#pooloptions) configuration.
+
+Note that specifying `scope: 'worker'` when running tests in `vmThreads` or `vmForks` will work the same way as `scope: 'file'`. This limitation exists because every test file has its own VM context, so if Vitest were to initiate it once, one context could leak to another and create many reference inconsistencies (instances of the same class would reference different constructors, for example).
+
 #### TypeScript
 
 To provide fixture types for all your custom contexts, you can pass the fixtures type as a generic.

packages/browser/src/client/public/esm-client-injector.js

@@ -20,6 +20,7 @@
     wrapModule,
     wrapDynamicImport: wrapModule,
     moduleCache,
+    cleanups: [],
     config: { __VITEST_CONFIG__ },
     viteConfig: { __VITEST_VITE_CONFIG__ },
     type: { __VITEST_TYPE__ },

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

@@ -1,5 +1,5 @@
-import type { CancelReason, File, Suite, Task, TaskEventPack, TaskResultPack, TestAnnotation, VitestRunner } from '@vitest/runner'
-import type { RunnerTestCase, SerializedConfig, TestExecutionMethod, WorkerGlobalState } from 'vitest'
+import type { CancelReason, File, Suite, Task, TaskEventPack, TaskResultPack, Test, TestAnnotation, VitestRunner } from '@vitest/runner'
+import type { SerializedConfig, TestExecutionMethod, WorkerGlobalState } from 'vitest'
 import type { VitestExecutor } from 'vitest/execute'
 import type { VitestBrowserClientMocker } from './mocker'
 import { globalChannel, onCancel } from '@vitest/browser/client'
@@ -59,7 +59,7 @@ export function createBrowserRunner(
       await super.onBeforeTryTask?.(...args)
     }
 
-    onAfterRunTask = async (task: Task) => {
+    onAfterRunTask = async (task: Test) => {
       await super.onAfterRunTask?.(task)
 
       if (this.config.bail && task.result?.state === 'fail') {
@@ -146,7 +146,7 @@ export function createBrowserRunner(
       return rpc().onCollected(this.method, files)
     }
 
-    onTestAnnotate = (test: RunnerTestCase, annotation: TestAnnotation): Promise<TestAnnotation> => {
+    onTestAnnotate = (test: Test, annotation: TestAnnotation): Promise<TestAnnotation> => {
       if (annotation.location) {
         // the file should be the test file
         // tests from other files are not supported

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

@@ -30,6 +30,7 @@ const state: WorkerGlobalState = {
       throw new Error('Not called in the browser')
     },
   },
+  onCleanup: fn => getBrowserState().cleanups.push(fn),
   moduleCache: getBrowserState().moduleCache,
   rpc: null as any,
   durations: {

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

@@ -234,6 +234,10 @@ async function cleanup() {
   await userEvent.cleanup()
     .catch(error => unhandledError(error, 'Cleanup Error'))
 
+  await Promise.all(
+    getBrowserState().cleanups.map(fn => fn()),
+  ).catch(error => unhandledError(error, 'Cleanup Error'))
+
   // if isolation is disabled, Vitest reuses the same iframe and we
   // don't need to switch the context back at all
   if (contextSwitched) {

packages/browser/src/client/utils.ts

@@ -81,6 +81,7 @@ export interface BrowserRunnerState {
   method: 'run' | 'collect'
   orchestrator?: IframeOrchestrator
   commands: CommandsManager
+  cleanups: Array<() => unknown>
   cdp?: {
     on: (event: string, listener: (payload: any) => void) => void
     once: (event: string, listener: (payload: any) => void) => void

packages/runner/src/context.ts

@@ -1,6 +1,7 @@
 import type { Awaitable } from '@vitest/utils'
 import type { VitestRunner } from './types/runner'
 import type {
+  File,
   RuntimeContext,
   SuiteCollector,
   Test,
@@ -271,6 +272,20 @@ function makeTimeoutError(isHook: boolean, timeout: number, stackTraceError?: Er
   return error
 }
 
+const fileContexts = new WeakMap<File, Record<string, unknown>>()
+
+export function getFileContext(file: File): Record<string, unknown> {
+  const context = fileContexts.get(file)
+  if (!context) {
+    throw new Error(`Cannot find file context for ${file.name}`)
+  }
+  return context
+}
+
+export function setFileContext(file: File, context: Record<string, unknown>): void {
+  fileContexts.set(file, context)
+}
+
 const table: string[] = []
 for (let i = 65; i < 91; i++) {
   table.push(String.fromCharCode(i))

packages/runner/src/fixture.ts

@@ -1,10 +1,13 @@
+import type { VitestRunner } from './types'
 import type { FixtureOptions, TestContext } from './types/tasks'
 import { createDefer, isObject } from '@vitest/utils'
+import { getFileContext } from './context'
 import { getTestFixture } from './map'
 
 export interface FixtureItem extends FixtureOptions {
   prop: string
   value: any
+  scope: 'test' | 'file' | 'worker'
   /**
    * Indicates whether the fixture is a function
    */
@@ -43,9 +46,9 @@ export function mergeScopedFixtures(
 export function mergeContextFixtures<T extends { fixtures?: FixtureItem[] }>(
   fixtures: Record<string, any>,
   context: T,
-  inject: (key: string) => unknown,
+  runner: VitestRunner,
 ): T {
-  const fixtureOptionKeys = ['auto', 'injected']
+  const fixtureOptionKeys = ['auto', 'injected', 'scope']
   const fixtureArray: FixtureItem[] = Object.entries(fixtures).map(
     ([prop, value]) => {
       const fixtureItem = { value } as FixtureItem
@@ -60,10 +63,14 @@ export function mergeContextFixtures<T extends { fixtures?: FixtureItem[] }>(
         Object.assign(fixtureItem, value[1])
         const userValue = value[0]
         fixtureItem.value = fixtureItem.injected
-          ? (inject(prop) ?? userValue)
+          ? (runner.injectValue?.(prop) ?? userValue)
           : userValue
       }
 
+      fixtureItem.scope = fixtureItem.scope || 'test'
+      if (fixtureItem.scope === 'worker' && !runner.getWorkerContext) {
+        fixtureItem.scope = 'file'
+      }
       fixtureItem.prop = prop
       fixtureItem.isFn = typeof fixtureItem.value === 'function'
       return fixtureItem
@@ -86,6 +93,25 @@ export function mergeContextFixtures<T extends { fixtures?: FixtureItem[] }>(
           ({ prop }) => prop !== fixture.prop && usedProps.includes(prop),
         )
       }
+      // test can access anything, so we ignore it
+      if (fixture.scope !== 'test') {
+        fixture.deps?.forEach((dep) => {
+          if (!dep.isFn) {
+            // non fn fixtures are always resolved and available to anyone
+            return
+          }
+          // worker scope can only import from worker scope
+          if (fixture.scope === 'worker' && dep.scope === 'worker') {
+            return
+          }
+          // file scope an import from file and worker scopes
+          if (fixture.scope === 'file' && dep.scope !== 'test') {
+            return
+          }
+
+          throw new SyntaxError(`cannot use the ${dep.scope} fixture "${dep.prop}" inside the ${fixture.scope} fixture "${fixture.prop}"`)
+        })
+      }
     }
   })
 
@@ -94,19 +120,19 @@ export function mergeContextFixtures<T extends { fixtures?: FixtureItem[] }>(
 
 const fixtureValueMaps = new Map<TestContext, Map<FixtureItem, any>>()
 const cleanupFnArrayMap = new Map<
-  TestContext,
+  object,
   Array<() => void | Promise<void>>
 >()
 
-export async function callFixtureCleanup(context: TestContext): Promise<void> {
+export async function callFixtureCleanup(context: object): Promise<void> {
   const cleanupFnArray = cleanupFnArrayMap.get(context) ?? []
   for (const cleanup of cleanupFnArray.reverse()) {
     await cleanup()
   }
   cleanupFnArrayMap.delete(context)
 }
 
-export function withFixtures(fn: Function, testContext?: TestContext) {
+export function withFixtures(runner: VitestRunner, fn: Function, testContext?: TestContext) {
   return (hookContext?: TestContext): any => {
     const context: (TestContext & { [key: string]: any }) | undefined
       = hookContext || testContext
@@ -153,21 +179,94 @@ export function withFixtures(fn: Function, testContext?: TestContext) {
           continue
         }
 
-        const resolvedValue = fixture.isFn
-          ? await resolveFixtureFunction(fixture.value, context, cleanupFnArray)
-          : fixture.value
+        const resolvedValue = await resolveFixtureValue(
+          runner,
+          fixture,
+          context!,
+          cleanupFnArray,
+        )
         context![fixture.prop] = resolvedValue
         fixtureValueMap.set(fixture, resolvedValue)
-        cleanupFnArray.unshift(() => {
-          fixtureValueMap.delete(fixture)
-        })
+
+        if (fixture.scope === 'test') {
+          cleanupFnArray.unshift(() => {
+            fixtureValueMap.delete(fixture)
+          })
+        }
       }
     }
 
     return resolveFixtures().then(() => fn(context))
   }
 }
 
+const globalFixturePromise = new WeakMap<FixtureItem, Promise<unknown>>()
+
+function resolveFixtureValue(
+  runner: VitestRunner,
+  fixture: FixtureItem,
+  context: TestContext & { [key: string]: any },
+  cleanupFnArray: (() => void | Promise<void>)[],
+) {
+  const fileContext = getFileContext(context.task.file)
+  const workerContext = runner.getWorkerContext?.()
+
+  if (!fixture.isFn) {
+    fileContext[fixture.prop] ??= fixture.value
+    if (workerContext) {
+      workerContext[fixture.prop] ??= fixture.value
+    }
+    return fixture.value
+  }
+
+  if (fixture.scope === 'test') {
+    return resolveFixtureFunction(
+      fixture.value,
+      context,
+      cleanupFnArray,
+    )
+  }
+
+  // in case the test runs in parallel
+  if (globalFixturePromise.has(fixture)) {
+    return globalFixturePromise.get(fixture)!
+  }
+
+  let fixtureContext: Record<string, unknown>
+
+  if (fixture.scope === 'worker') {
+    if (!workerContext) {
+      throw new TypeError('[@vitest/runner] The worker context is not available in the current test runner. Please, provide the `getWorkerContext` method when initiating the runner.')
+    }
+    fixtureContext = workerContext
+  }
+  else {
+    fixtureContext = fileContext
+  }
+
+  if (fixture.prop in fixtureContext) {
+    return fixtureContext[fixture.prop]
+  }
+
+  if (!cleanupFnArrayMap.has(fixtureContext)) {
+    cleanupFnArrayMap.set(fixtureContext, [])
+  }
+  const cleanupFnFileArray = cleanupFnArrayMap.get(fixtureContext)!
+
+  const promise = resolveFixtureFunction(
+    fixture.value,
+    fixtureContext,
+    cleanupFnFileArray,
+  ).then((value) => {
+    fixtureContext[fixture.prop] = value
+    globalFixturePromise.delete(fixture)
+    return value
+  })
+
+  globalFixturePromise.set(fixture, promise)
+  return promise
+}
+
 async function resolveFixtureFunction(
   fixtureFn: (
     context: unknown,

packages/runner/src/hooks.ts

@@ -139,11 +139,12 @@ export function beforeEach<ExtraContext = object>(
 ): void {
   assertTypes(fn, '"beforeEach" callback', ['function'])
   const stackTraceError = new Error('STACK_TRACE_ERROR')
+  const runner = getRunner()
   return getCurrentSuite<ExtraContext>().on(
     'beforeEach',
     Object.assign(
       withTimeout(
-        withFixtures(fn),
+        withFixtures(runner, fn),
         timeout ?? getDefaultHookTimeout(),
         true,
         stackTraceError,
@@ -179,10 +180,11 @@ export function afterEach<ExtraContext = object>(
   timeout?: number,
 ): void {
   assertTypes(fn, '"afterEach" callback', ['function'])
+  const runner = getRunner()
   return getCurrentSuite<ExtraContext>().on(
     'afterEach',
     withTimeout(
-      withFixtures(fn),
+      withFixtures(runner, fn),
       timeout ?? getDefaultHookTimeout(),
       true,
       new Error('STACK_TRACE_ERROR'),