Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: implement provide/inject API to transfer data from the main thread #4422

Merged
merged 10 commits into from
Nov 4, 2023
25 changes: 24 additions & 1 deletion docs/config/index.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
}
```
:::


Expand Down
1 change: 1 addition & 0 deletions packages/browser/src/client/main.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
4 changes: 4 additions & 0 deletions packages/vitest/src/api/setup.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down
3 changes: 2 additions & 1 deletion packages/vitest/src/api/types.ts
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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'> {
Expand Down
1 change: 1 addition & 0 deletions packages/vitest/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down
11 changes: 11 additions & 0 deletions packages/vitest/src/integrations/inject.ts
Original file line number Diff line number Diff line change
@@ -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]
}
25 changes: 5 additions & 20 deletions packages/vitest/src/integrations/vi.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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(
Expand Down Expand Up @@ -587,8 +576,7 @@ function createVitest(): VitestUtils {
},

resetModules() {
const state = getWorkerState()
resetModules(state.moduleCache)
resetModules(workerState.moduleCache)
return utils
},

Expand All @@ -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)
},
}

Expand Down
6 changes: 6 additions & 0 deletions packages/vitest/src/node/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [
Expand Down
41 changes: 23 additions & 18 deletions packages/vitest/src/node/core.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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[]) {
Expand Down Expand Up @@ -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
}
Expand Down
9 changes: 8 additions & 1 deletion packages/vitest/src/node/globalSetup.ts
Original file line number Diff line number Diff line change
@@ -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
}

Expand Down
1 change: 1 addition & 0 deletions packages/vitest/src/node/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down
1 change: 1 addition & 0 deletions packages/vitest/src/node/pools/child.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 })
Expand Down
1 change: 1 addition & 0 deletions packages/vitest/src/node/pools/threads.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 })
Expand Down
1 change: 1 addition & 0 deletions packages/vitest/src/node/pools/vm-threads.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 })
Expand Down
39 changes: 30 additions & 9 deletions packages/vitest/src/node/workspace.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down Expand Up @@ -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,
Expand All @@ -85,17 +85,38 @@ export class WorkspaceProject {
return this.ctx.getCoreWorkspaceProject() === this
}

provide = (key: string, value: unknown) => {
sheremet-va marked this conversation as resolved.
Show resolved Hide resolved
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 })
sheremet-va marked this conversation as resolved.
Show resolved Hide resolved
if (teardown == null || !!globalSetupFile.teardown)
continue
if (typeof teardown !== 'function')
Expand All @@ -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 {
Expand Down Expand Up @@ -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
Expand Down
3 changes: 2 additions & 1 deletion packages/vitest/src/runtime/child.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -72,6 +72,7 @@ async function init(ctx: ChildContext) {
prepare: performance.now(),
},
rpc,
providedContext,
isChildProcess: true,
}

Expand Down
7 changes: 4 additions & 3 deletions packages/vitest/src/runtime/execute.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down Expand Up @@ -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
AriPerkkio marked this conversation as resolved.
Show resolved Hide resolved
const rpc = () => state().rpc

const processExit = process.exit
Expand Down Expand Up @@ -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 {
Expand Down
Loading
Loading