feat: implement provide/inject API to transfer data from the main thr…

…ead (#4422)

docs/config/index.md

@@ -957,7 +957,30 @@ Multiple globalSetup files are possible. setup and teardown are executed sequent
 :::
 
 ::: warning
-Beware that the global setup is running in a different global scope, so your tests don't have access to variables defined here. Also, since Vitest 1.0.0-beta, global setup runs only if there is at least one running test. This means that global setup might start running during watch mode after test file is changed, for example (the test file will wait for global setup to finish before running).
+Since Vitest 1.0.0-beta, global setup runs only if there is at least one running test. This means that global setup might start running during watch mode after test file is changed (the test file will wait for global setup to finish before running).
+
+Beware that the global setup is running in a different global scope, so your tests don't have access to variables defined here. Hovewer, since 1.0.0 you can pass down serializable data to tests via `provide` method:
+
+```ts
+// globalSetup.js
+export default function setup({ provide }) {
+  provide('wsPort', 3000)
+}
+// example.test.js
+import { inject } from 'vitest'
+
+inject('wsPort') === 3000
+```
+
+If you are using TypeScript, you can extend `ProvidedContext` type to have type safe access to `provide/inject` methods:
+
+```ts
+declare module 'vitest' {
+  export interface ProvidedContext {
+    wsPort: number
+  }
+}
+```
 :::
 
 

packages/browser/src/client/main.ts

@@ -120,6 +120,7 @@ ws.addEventListener('open', async () => {
       environment: 0,
       prepare: 0,
     },
+    providedContext: await client.rpc.getProvidedContext(),
   }
   // @ts-expect-error mocking vitest apis
   globalThis.__vitest_mocker__ = new VitestBrowserClientMocker()

packages/vitest/src/api/setup.ts

@@ -131,6 +131,10 @@ export function setup(vitestOrWorkspace: Vitest | WorkspaceProject, server?: Vit
         getCountOfFailedTests() {
           return ctx.state.getCountOfFailedTests()
         },
+        // browser should have a separate RPC in the future, UI doesn't care for provided context
+        getProvidedContext() {
+          return 'ctx' in vitestOrWorkspace ? vitestOrWorkspace.getProvidedContext() : ({} as any)
+        },
       },
       {
         post: msg => ws.send(msg),

packages/vitest/src/api/types.ts

@@ -1,6 +1,6 @@
 import type { TransformResult } from 'vite'
 import type { CancelReason } from '@vitest/runner'
-import type { AfterSuiteRunMeta, File, ModuleGraphData, Reporter, ResolvedConfig, SnapshotResult, TaskResultPack, UserConsoleLog } from '../types'
+import type { AfterSuiteRunMeta, File, ModuleGraphData, ProvidedContext, Reporter, ResolvedConfig, SnapshotResult, TaskResultPack, UserConsoleLog } from '../types'
 
 export interface TransformResultWithSource extends TransformResult {
   source?: string
@@ -30,6 +30,7 @@ export interface WebSocketHandlers {
   snapshotSaved(snapshot: SnapshotResult): void
   rerun(files: string[]): Promise<void>
   updateSnapshot(file?: File): Promise<void>
+  getProvidedContext(): ProvidedContext
 }
 
 export interface WebSocketEvents extends Pick<Reporter, 'onCollected' | 'onFinished' | 'onTaskUpdate' | 'onUserConsoleLog' | 'onPathsCollected'> {

packages/vitest/src/index.ts

@@ -15,6 +15,7 @@ export { runOnce, isFirstRun } from './integrations/run-once'
 export * from './integrations/chai'
 export * from './integrations/vi'
 export * from './integrations/utils'
+export { inject } from './integrations/inject'
 export type { SnapshotEnvironment } from '@vitest/snapshot/environment'
 
 export * from './types'

packages/vitest/src/integrations/inject.ts

@@ -0,0 +1,11 @@
+import type { ProvidedContext } from '../types/general'
+import { getWorkerState } from '../utils/global'
+
+/**
+ * Gives access to injected context provided from the main thread.
+ * This usually returns a value provided by `globalSetup` or an external library.
+ */
+export function inject<T extends keyof ProvidedContext>(key: T): ProvidedContext[T] {
+  const workerState = getWorkerState()
+  return workerState.providedContext[key]
+}

packages/vitest/src/integrations/vi.ts

@@ -353,15 +353,6 @@ function createVitest(): VitestUtils {
 
   const workerState = getWorkerState()
 
-  if (!workerState) {
-    const errorMsg = 'Vitest failed to access its internal state.'
-      + '\n\nOne of the following is possible:'
-      + '\n- "vitest" is imported directly without running "vitest" command'
-      + '\n- "vitest" is imported inside "globalSetup" (to fix this, use "setupFiles" instead, because "globalSetup" runs in a different context)'
-      + '\n- Otherwise, it might be a Vitest bug. Please report it to https://github.com/vitest-dev/vitest/issues\n'
-    throw new Error(errorMsg)
-  }
-
   const _timers = new FakeTimers({
     global: globalThis,
     config: workerState.config.fakeTimers,
@@ -379,8 +370,6 @@ function createVitest(): VitestUtils {
 
   const utils: VitestUtils = {
     useFakeTimers(config?: FakeTimerInstallOpts) {
-      const workerState = getWorkerState()
-
       if (workerState.isChildProcess) {
         if (config?.toFake?.includes('nextTick') || workerState.config?.fakeTimers?.toFake?.includes('nextTick')) {
           throw new Error(
@@ -587,8 +576,7 @@ function createVitest(): VitestUtils {
     },
 
     resetModules() {
-      const state = getWorkerState()
-      resetModules(state.moduleCache)
+      resetModules(workerState.moduleCache)
       return utils
     },
 
@@ -597,17 +585,14 @@ function createVitest(): VitestUtils {
     },
 
     setConfig(config: RuntimeConfig) {
-      const state = getWorkerState()
       if (!_config)
-        _config = { ...state.config }
-      Object.assign(state.config, config)
+        _config = { ...workerState.config }
+      Object.assign(workerState.config, config)
     },
 
     resetConfig() {
-      if (_config) {
-        const state = getWorkerState()
-        Object.assign(state.config, _config)
-      }
+      if (_config)
+        Object.assign(workerState.config, _config)
     },
   }
 

packages/vitest/src/node/config.ts

@@ -301,6 +301,12 @@ export function resolveConfig(
         ?? resolve(resolved.root, file),
     ),
   )
+  resolved.globalSetup = toArray(resolved.globalSetup || []).map(file =>
+    normalize(
+      resolveModule(file, { paths: [resolved.root] })
+        ?? resolve(resolved.root, file),
+    ),
+  )
   resolved.coverage.exclude.push(...resolved.setupFiles.map(file => `${resolved.coverage.allowExternal ? '**/' : ''}${relative(resolved.root, file)}`))
 
   resolved.forceRerunTriggers = [

packages/vitest/src/node/core.ts

@@ -438,9 +438,8 @@ export class Vitest {
     const coreProject = this.getCoreWorkspaceProject()
     if (!projects.has(coreProject))
       projects.add(coreProject)
-    await Promise.all(
-      Array.from(projects).map(project => project.initializeGlobalSetup()),
-    )
+    for await (const project of projects)
+      await project.initializeGlobalSetup()
   }
 
   async runFiles(paths: WorkspaceSpec[]) {
@@ -750,22 +749,28 @@ export class Vitest {
 
   async close() {
     if (!this.closingPromise) {
-      const closePromises: unknown[] = this.projects.map(w => w.close().then(() => w.server = undefined as any))
-      // close the core workspace server only once
-      // it's possible that it's not initialized at all because it's not running any tests
-      if (!this.projects.includes(this.coreWorkspaceProject))
-        closePromises.push(this.coreWorkspaceProject.close().then(() => this.server = undefined as any))
-
-      if (this.pool)
-        closePromises.push(this.pool.close().then(() => this.pool = undefined))
-
-      closePromises.push(...this._onClose.map(fn => fn()))
-
-      this.closingPromise = Promise.allSettled(closePromises).then((results) => {
-        results.filter(r => r.status === 'rejected').forEach((err) => {
-          this.logger.error('error during close', (err as PromiseRejectedResult).reason)
+      this.closingPromise = (async () => {
+        // do teardown before closing the server
+        for await (const project of [...this.projects].reverse())
+          await project.teardownGlobalSetup()
+
+        const closePromises: unknown[] = this.projects.map(w => w.close().then(() => w.server = undefined as any))
+        // close the core workspace server only once
+        // it's possible that it's not initialized at all because it's not running any tests
+        if (!this.projects.includes(this.coreWorkspaceProject))
+          closePromises.push(this.coreWorkspaceProject.close().then(() => this.server = undefined as any))
+
+        if (this.pool)
+          closePromises.push(this.pool.close().then(() => this.pool = undefined))
+
+        closePromises.push(...this._onClose.map(fn => fn()))
+
+        return Promise.allSettled(closePromises).then((results) => {
+          results.filter(r => r.status === 'rejected').forEach((err) => {
+            this.logger.error('error during close', (err as PromiseRejectedResult).reason)
+          })
         })
-      })
+      })()
     }
     return this.closingPromise
   }

packages/vitest/src/node/globalSetup.ts

@@ -1,9 +1,16 @@
 import { toArray } from '@vitest/utils'
 import type { ViteNodeRunner } from 'vite-node/client'
+import type { ProvidedContext } from '../types/general'
+import type { ResolvedConfig } from '../types/config'
+
+export interface GlobalSetupContext {
+  config: ResolvedConfig
+  provide<T extends keyof ProvidedContext>(key: T, value: ProvidedContext[T]): void
+}
 
 export interface GlobalSetupFile {
   file: string
-  setup?: () => Promise<Function | void> | void
+  setup?: (context: GlobalSetupContext) => Promise<Function | void> | void
   teardown?: Function
 }
 

packages/vitest/src/node/index.ts

@@ -5,6 +5,7 @@ export { VitestPlugin } from './plugins'
 export { startVitest } from './cli-api'
 export { registerConsoleShortcuts } from './stdin'
 export type { WorkspaceSpec } from './pool'
+export type { GlobalSetupContext } from './globalSetup'
 
 export type { TestSequencer, TestSequencerConstructor } from './sequencers/types'
 export { BaseSequencer } from './sequencers/BaseSequencer'

packages/vitest/src/node/pools/child.ts

@@ -104,6 +104,7 @@ export function createChildProcessPool(ctx: Vitest, { execArgv, env, forksPath }
         environment,
         workerId,
         projectName: project.getName(),
+        providedContext: project.getProvidedContext(),
       }
       try {
         await pool.run(data, { name, channel })

packages/vitest/src/node/pools/threads.ts

@@ -92,6 +92,7 @@ export function createThreadsPool(ctx: Vitest, { execArgv, env, workerPath }: Po
         environment,
         workerId,
         projectName: project.getName(),
+        providedContext: project.getProvidedContext(),
       }
       try {
         await pool.run(data, { transferList: [workerPort], name })

packages/vitest/src/node/pools/vm-threads.ts

@@ -97,6 +97,7 @@ export function createVmThreadsPool(ctx: Vitest, { execArgv, env, vmPath }: Pool
         environment,
         workerId,
         projectName: project.getName(),
+        providedContext: project.getProvidedContext(),
       }
       try {
         await pool.run(data, { transferList: [workerPort], name })

packages/vitest/src/node/workspace.ts

@@ -8,7 +8,7 @@ import { ViteNodeServer } from 'vite-node/server'
 import c from 'picocolors'
 import type { RawSourceMap } from 'vite-node'
 import { createBrowserServer } from '../integrations/browser/server'
-import type { ArgumentsType, Reporter, ResolvedConfig, UserConfig, UserWorkspaceConfig, Vitest } from '../types'
+import type { ArgumentsType, ProvidedContext, Reporter, ResolvedConfig, UserConfig, UserWorkspaceConfig, Vitest } from '../types'
 import { deepMerge } from '../utils'
 import type { Typechecker } from '../typecheck/typechecker'
 import type { BrowserProvider } from '../types/browser'
@@ -69,8 +69,8 @@ export class WorkspaceProject {
 
   testFilesList: string[] = []
 
-  private _globalSetupInit = false
-  private _globalSetups: GlobalSetupFile[] = []
+  private _globalSetups: GlobalSetupFile[] | undefined
+  private _provided: ProvidedContext = {} as any
 
   constructor(
     public path: string | number,
@@ -85,17 +85,38 @@ export class WorkspaceProject {
     return this.ctx.getCoreWorkspaceProject() === this
   }
 
+  provide = (key: string, value: unknown) => {
+    try {
+      structuredClone(value)
+    }
+    catch (err) {
+      throw new Error(`Cannot provide "${key}" because it's not serializable.`, {
+        cause: err,
+      })
+    }
+    (this._provided as any)[key] = value
+  }
+
+  getProvidedContext(): ProvidedContext {
+    if (this.isCore())
+      return this._provided
+    // globalSetup can run even if core workspace is not part of the test run
+    // so we need to inherit its provided context
+    return {
+      ...this.ctx.getCoreWorkspaceProject().getProvidedContext(),
+      ...this._provided,
+    }
+  }
+
   async initializeGlobalSetup() {
-    if (this._globalSetupInit)
+    if (this._globalSetups)
       return
 
-    this._globalSetupInit = true
-
     this._globalSetups = await loadGlobalSetupFiles(this.runner, this.config.globalSetup)
 
     try {
       for (const globalSetupFile of this._globalSetups) {
-        const teardown = await globalSetupFile.setup?.()
+        const teardown = await globalSetupFile.setup?.({ provide: this.provide, config: this.config })
         if (teardown == null || !!globalSetupFile.teardown)
           continue
         if (typeof teardown !== 'function')
@@ -111,7 +132,7 @@ export class WorkspaceProject {
   }
 
   async teardownGlobalSetup() {
-    if (!this._globalSetupInit || !this._globalSetups.length)
+    if (!this._globalSetups)
       return
     for (const globalSetupFile of [...this._globalSetups].reverse()) {
       try {
@@ -338,7 +359,7 @@ export class WorkspaceProject {
         this.server.close(),
         this.typechecker?.stop(),
         this.browser?.close(),
-        this.teardownGlobalSetup(),
+        () => this._provided = ({} as any),
       ].filter(Boolean))
     }
     return this.closingPromise

packages/vitest/src/runtime/child.ts

@@ -14,7 +14,7 @@ import { createSafeRpc, rpcDone } from './rpc'
 import { setupInspect } from './inspector'
 
 async function init(ctx: ChildContext) {
-  const { config, workerId } = ctx
+  const { config, workerId, providedContext } = ctx
 
   process.env.VITEST_WORKER_ID = String(workerId)
   process.env.VITEST_POOL_ID = String(poolId)
@@ -72,6 +72,7 @@ async function init(ctx: ChildContext) {
       prepare: performance.now(),
     },
     rpc,
+    providedContext,
     isChildProcess: true,
   }
 

packages/vitest/src/runtime/execute.ts

@@ -8,7 +8,6 @@ import { processError } from '@vitest/utils/error'
 import type { MockMap } from '../types/mocker'
 import type { ResolvedConfig, ResolvedTestEnvironment, RuntimeRPC, WorkerGlobalState } from '../types'
 import { distDir } from '../paths'
-import { getWorkerState } from '../utils/global'
 import { VitestMocker } from './mocker'
 import { ExternalModulesExecutor } from './external-executor'
 import { FileMap } from './vm/file-map'
@@ -64,7 +63,8 @@ export interface ContextExecutorOptions {
 }
 
 export async function startVitestExecutor(options: ContextExecutorOptions) {
-  const state = () => getWorkerState() || options.state
+  // @ts-expect-error injected untyped global
+  const state = () => globalThis.__vitest_worker__ || options.state
   const rpc = () => state().rpc
 
   const processExit = process.exit
@@ -203,7 +203,8 @@ export class VitestExecutor extends ViteNodeRunner {
   }
 
   get state() {
-    return getWorkerState() || this.options.state
+    // @ts-expect-error injected untyped global
+    return globalThis.__vitest_worker__ || this.options.state
   }
 
   shouldResolveId(id: string, _importee?: string | undefined): boolean {