Skip to content

Commit

Permalink
feat!: separate environments always run in a separate worker
Browse files Browse the repository at this point in the history
  • Loading branch information
sheremet-va committed Feb 21, 2023
1 parent 8ab33f1 commit 4e53b84
Show file tree
Hide file tree
Showing 6 changed files with 120 additions and 94 deletions.
37 changes: 31 additions & 6 deletions packages/vitest/src/node/pool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,10 @@ import type { Options as TinypoolOptions } from 'tinypool'
import { Tinypool } from 'tinypool'
import { createBirpc } from 'birpc'
import type { RawSourceMap } from 'vite-node'
import type { ResolvedConfig, WorkerContext, WorkerRPC } from '../types'
import type { ResolvedConfig, WorkerContext, WorkerRPC, WorkerTestEnvironment } from '../types'
import { distDir, rootDir } from '../constants'
import { AggregateError } from '../utils'
import { AggregateError, groupBy } from '../utils'
import { envsOrder, groupFilesByEnv } from '../utils/test-helpers'
import type { Vitest } from './core'

export type RunWithFiles = (files: string[], invalidates?: string[]) => Promise<void>
Expand Down Expand Up @@ -88,7 +89,7 @@ export function createPool(ctx: Vitest): WorkerPool {
const runWithFiles = (name: string): RunWithFiles => {
let id = 0

async function runFiles(config: ResolvedConfig, files: string[], invalidates: string[] = []) {
async function runFiles(config: ResolvedConfig, files: string[], environment: WorkerTestEnvironment, invalidates: string[] = []) {
ctx.state.clearFiles(files)
const { workerPort, port } = createChannel(ctx)
const workerId = ++id
Expand All @@ -97,6 +98,7 @@ export function createPool(ctx: Vitest): WorkerPool {
config,
files,
invalidates,
environment,
workerId,
}
try {
Expand All @@ -119,12 +121,35 @@ export function createPool(ctx: Vitest): WorkerPool {

files = await sequencer.sort(files)

const filesByEnv = await groupFilesByEnv(files, config)
const envs = envsOrder.concat(
Object.keys(filesByEnv).filter(env => !envsOrder.includes(env)),
)

if (!ctx.config.threads) {
await runFiles(config, files)
// always run environments isolated between each other
for (const env of envs) {
const files = filesByEnv[env]

if (!files?.length)
continue

const filesByOptions = groupBy(files, ({ environment }) => JSON.stringify(environment.options))

for (const option in filesByOptions) {
const files = filesByOptions[option]

if (files?.length) {
const filenames = files.map(f => f.file)
await runFiles(config, filenames, files[0].environment, invalidates)
}
}
}
}
else {
const results = await Promise.allSettled(files
.map(file => runFiles(config, [file], invalidates)))
const promises = Object.values(filesByEnv).flat()
const results = await Promise.allSettled(promises
.map(({ file, environment }) => runFiles(config, [file], environment, invalidates)))

const errors = results.filter((r): r is PromiseRejectedResult => r.status === 'rejected').map(r => r.reason)
if (errors.length > 0)
Expand Down
106 changes: 22 additions & 84 deletions packages/vitest/src/runtime/entry.ts
Original file line number Diff line number Diff line change
@@ -1,12 +1,9 @@
import { promises as fs } from 'node:fs'
import mm from 'micromatch'
import type { VitestRunner, VitestRunnerConstructor } from '@vitest/runner'
import { startTests } from '@vitest/runner'
import { resolve } from 'pathe'
import type { EnvironmentOptions, ResolvedConfig, VitestEnvironment } from '../types'
import type { ResolvedConfig, WorkerTestEnvironment } from '../types'
import { getWorkerState, resetModules } from '../utils'
import { vi } from '../integrations/vi'
import { envs } from '../integrations/env'
import { distDir } from '../constants'
import { startCoverageInsideWorker, stopCoverageInsideWorker, takeCoverageInsideWorker } from '../integrations/coverage'
import { setupGlobalEnv, withEnv } from './setup.node'
Expand All @@ -15,15 +12,6 @@ import type { VitestExecutor } from './execute'

const runnersFile = resolve(distDir, 'runners.js')

function groupBy<T, K extends string | number | symbol>(collection: T[], iteratee: (item: T) => K) {
return collection.reduce((acc, item) => {
const key = iteratee(item)
acc[key] ||= []
acc[key].push(item)
return acc
}, {} as Record<K, T[]>)
}

async function getTestRunnerConstructor(config: ResolvedConfig, executor: VitestExecutor): Promise<VitestRunnerConstructor> {
if (!config.runner) {
const { VitestTestRunner, NodeBenchmarkRunner } = await executor.executeFile(runnersFile)
Expand Down Expand Up @@ -77,89 +65,39 @@ async function getTestRunner(config: ResolvedConfig, executor: VitestExecutor):
}

// browser shouldn't call this!
export async function run(files: string[], config: ResolvedConfig, executor: VitestExecutor): Promise<void> {
export async function run(files: string[], config: ResolvedConfig, environment: WorkerTestEnvironment, executor: VitestExecutor): Promise<void> {
await setupGlobalEnv(config)
await startCoverageInsideWorker(config.coverage, executor)

const workerState = getWorkerState()

const runner = await getTestRunner(config, executor)

// if calling from a worker, there will always be one file
// if calling with no-threads, this will be the whole suite
const filesWithEnv = await Promise.all(files.map(async (file) => {
const code = await fs.readFile(file, 'utf-8')

// 1. Check for control comments in the file
let env = code.match(/@(?:vitest|jest)-environment\s+?([\w-]+)\b/)?.[1]
// 2. Check for globals
if (!env) {
for (const [glob, target] of config.environmentMatchGlobs || []) {
if (mm.isMatch(file, glob)) {
env = target
break
}
// @ts-expect-error untyped global
globalThis.__vitest_environment__ = environment

await withEnv(environment.name, environment.options || config.environmentOptions || {}, executor, async () => {
for (const file of files) {
// it doesn't matter if running with --threads
// if running with --no-threads, we usually want to reset everything before running a test
// but we have --isolate option to disable this
if (config.isolate) {
workerState.mockMap.clear()
resetModules(workerState.moduleCache, true)
}
}
// 3. Fallback to global env
env ||= config.environment || 'node'

const envOptions = JSON.parse(code.match(/@(?:vitest|jest)-environment-options\s+?(.+)/)?.[1] || 'null')
return {
file,
env: env as VitestEnvironment,
envOptions: envOptions ? { [env]: envOptions } as EnvironmentOptions : null,
}
}))

const filesByEnv = groupBy(filesWithEnv, ({ env }) => env)

const orderedEnvs = envs.concat(
Object.keys(filesByEnv).filter(env => !envs.includes(env)),
)

for (const env of orderedEnvs) {
const environment = env as VitestEnvironment
const files = filesByEnv[environment]

if (!files || !files.length)
continue

// @ts-expect-error untyped global
globalThis.__vitest_environment__ = environment
workerState.filepath = file

const filesByOptions = groupBy(files, ({ envOptions }) => JSON.stringify(envOptions))
await startTests([file], runner)

for (const options of Object.keys(filesByOptions)) {
const files = filesByOptions[options]
workerState.filepath = undefined

if (!files || !files.length)
continue

await withEnv(environment, files[0].envOptions || config.environmentOptions || {}, executor, async () => {
for (const { file } of files) {
// it doesn't matter if running with --threads
// if running with --no-threads, we usually want to reset everything before running a test
// but we have --isolate option to disable this
if (config.isolate) {
workerState.mockMap.clear()
resetModules(workerState.moduleCache, true)
}

workerState.filepath = file

await startTests([file], runner)

workerState.filepath = undefined

// reset after tests, because user might call `vi.setConfig` in setupFile
vi.resetConfig()
// mocks should not affect different files
vi.restoreAllMocks()
}
})
// reset after tests, because user might call `vi.setConfig` in setupFile
vi.resetConfig()
// mocks should not affect different files
vi.restoreAllMocks()
}
}

await stopCoverageInsideWorker(config.coverage, executor)
await stopCoverageInsideWorker(config.coverage, executor)
})
}
6 changes: 3 additions & 3 deletions packages/vitest/src/runtime/worker.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ import { workerId as poolId } from 'tinypool'
import { processError } from '@vitest/runner/utils'
import { ModuleCacheMap } from 'vite-node/client'
import { isPrimitive } from 'vite-node/utils'
import type { ResolvedConfig, WorkerContext, WorkerRPC } from '../types'
import type { ResolvedConfig, WorkerContext, WorkerRPC, WorkerTestEnvironment } from '../types'
import { distDir } from '../constants'
import { getWorkerState } from '../utils/global'
import type { MockMap } from '../types/mocker'
Expand All @@ -14,7 +14,7 @@ import { createVitestExecutor } from './execute'
import { rpc, rpcDone } from './rpc'

let _viteNode: {
run: (files: string[], config: ResolvedConfig, executor: VitestExecutor) => Promise<void>
run: (files: string[], config: ResolvedConfig, environment: WorkerTestEnvironment, executor: VitestExecutor) => Promise<void>
executor: VitestExecutor
}

Expand Down Expand Up @@ -109,6 +109,6 @@ function init(ctx: WorkerContext) {
export async function run(ctx: WorkerContext) {
init(ctx)
const { run, executor } = await startViteNode(ctx)
await run(ctx.files, ctx.config, executor)
await run(ctx.files, ctx.config, ctx.environment, executor)
await rpcDone()
}
8 changes: 7 additions & 1 deletion packages/vitest/src/types/worker.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,15 +3,21 @@ import type { File, TaskResultPack, Test } from '@vitest/runner'
import type { FetchFunction, ModuleCacheMap, RawSourceMap, ViteNodeResolveId } from 'vite-node'
import type { BirpcReturn } from 'birpc'
import type { MockMap } from './mocker'
import type { ResolvedConfig } from './config'
import type { EnvironmentOptions, ResolvedConfig, VitestEnvironment } from './config'
import type { SnapshotResult } from './snapshot'
import type { UserConsoleLog } from './general'

export interface WorkerTestEnvironment {
name: VitestEnvironment
options: EnvironmentOptions | null
}

export interface WorkerContext {
workerId: number
port: MessagePort
config: ResolvedConfig
files: string[]
environment: WorkerTestEnvironment
invalidates?: string[]
}

Expand Down
9 changes: 9 additions & 0 deletions packages/vitest/src/utils/base.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,15 @@ function collectOwnProperties(obj: any, collector: Set<string | symbol> | ((key:
Object.getOwnPropertySymbols(obj).forEach(collect)
}

export function groupBy<T, K extends string | number | symbol>(collection: T[], iteratee: (item: T) => K) {
return collection.reduce((acc, item) => {
const key = iteratee(item)
acc[key] ||= []
acc[key].push(item)
return acc
}, {} as Record<K, T[]>)
}

export function getAllMockableProperties(obj: any, isModule: boolean) {
const allProps = new Map<string | symbol, { key: string | symbol; descriptor: PropertyDescriptor }>()
let curr = obj
Expand Down
48 changes: 48 additions & 0 deletions packages/vitest/src/utils/test-helpers.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
import { promises as fs } from 'node:fs'
import mm from 'micromatch'
import type { EnvironmentOptions, ResolvedConfig, VitestEnvironment } from '../types'
import { groupBy } from './base'

export const envsOrder = [
'node',
'jsdom',
'happy-dom',
'edge-runtime',
]

export interface FileByEnv {
file: string
env: VitestEnvironment
envOptions: EnvironmentOptions | null
}

export async function groupFilesByEnv(files: string[], config: ResolvedConfig) {
const filesWithEnv = await Promise.all(files.map(async (file) => {
const code = await fs.readFile(file, 'utf-8')

// 1. Check for control comments in the file
let env = code.match(/@(?:vitest|jest)-environment\s+?([\w-]+)\b/)?.[1]
// 2. Check for globals
if (!env) {
for (const [glob, target] of config.environmentMatchGlobs || []) {
if (mm.isMatch(file, glob)) {
env = target
break
}
}
}
// 3. Fallback to global env
env ||= config.environment || 'node'

const envOptions = JSON.parse(code.match(/@(?:vitest|jest)-environment-options\s+?(.+)/)?.[1] || 'null')
return {
file,
environment: {
name: env as VitestEnvironment,
options: envOptions ? { [env]: envOptions } as EnvironmentOptions : null,
},
}
}))

return groupBy(filesWithEnv, ({ environment }) => environment.name)
}

0 comments on commit 4e53b84

Please sign in to comment.