feat: web worker support (#726)

Co-authored-by: Anthony Fu <anthonyfu117@hotmail.com>

packages/vitest/src/integrations/jest-mock.ts

@@ -1,4 +1,3 @@
-import { util } from 'chai'
 import type { SpyImpl } from 'tinyspy'
 import * as tinyspy from 'tinyspy'
 
@@ -228,7 +227,9 @@ function enhanceSpy<TArgs extends any[], TReturns>(
   stub.mockRejectedValueOnce = (val: unknown) =>
     stub.mockImplementationOnce(() => Promise.reject(val))
 
-  util.addProperty(stub, 'mock', () => mockContext)
+  Object.defineProperty(stub, 'mock', {
+    get: () => mockContext,
+  })
 
   stub.willCall(function(this: unknown, ...args) {
     instances.push(this)

packages/vitest/src/integrations/snapshot/client.ts

@@ -2,7 +2,7 @@ import path from 'pathe'
 import { expect } from 'chai'
 import type { SnapshotResult, Test } from '../../types'
 import { rpc } from '../../runtime/rpc'
-import { getNames } from '../../utils'
+import { getNames, getWorkerState } from '../../utils'
 import { equals, iterableEquality, subsetEquality } from '../chai/jest-utils'
 import { deepMergeSnapshot } from './port/utils'
 import SnapshotState from './port/state'
@@ -34,7 +34,7 @@ export class SnapshotClient {
       this.testFile = this.test!.file!.filepath
       this.snapshotState = new SnapshotState(
         resolveSnapshotPath(this.testFile),
-        __vitest_worker__!.config.snapshotOptions,
+        getWorkerState().config.snapshotOptions,
       )
     }
   }

packages/vitest/src/integrations/vi.ts

@@ -1,7 +1,7 @@
 /* eslint-disable @typescript-eslint/no-unused-vars */
 
 import { parseStacktrace } from '../utils/source-map'
-import type { VitestMocker } from '../node/mocker'
+import type { VitestMocker } from '../runtime/mocker'
 import { FakeTimers } from './timers'
 import type { EnhancedSpy, MaybeMocked, MaybeMockedDeep } from './jest-mock'
 import { fn, isMockFunction, spies, spyOn } from './jest-mock'

packages/vitest/src/node/index.ts

@@ -2,3 +2,6 @@ export type { Vitest } from './core'
 export { createVitest } from './create'
 export { VitestPlugin } from './plugins'
 export { startVitest } from './cli-api'
+
+export { VitestRunner } from '../runtime/execute'
+export type { ExecuteOptions } from '../runtime/execute'

packages/vitest/src/runtime/context.ts

@@ -1,4 +1,5 @@
 import type { Awaitable, DoneCallback, RuntimeContext, SuiteCollector, TestFunction } from '../types'
+import { getWorkerState } from '../utils'
 
 export const context: RuntimeContext = {
   tasks: [],
@@ -17,11 +18,11 @@ export async function runWithSuite(suite: SuiteCollector, fn: (() => Awaitable<v
 }
 
 export function getDefaultTestTimeout() {
-  return __vitest_worker__!.config!.testTimeout
+  return getWorkerState().config.testTimeout
 }
 
 export function getDefaultHookTimeout() {
-  return __vitest_worker__!.config!.hookTimeout
+  return getWorkerState().config.hookTimeout
 }
 
 export function withTimeout<T extends((...args: any[]) => any)>(

packages/vitest/src/runtime/entry.ts

@@ -1,5 +1,6 @@
 import { promises as fs } from 'fs'
 import type { BuiltinEnvironment, ResolvedConfig } from '../types'
+import { getWorkerState } from '../utils'
 import { setupGlobalEnv, withEnv } from './setup'
 import { startTests } from './run'
 
@@ -14,12 +15,13 @@ export async function run(files: string[], config: ResolvedConfig): Promise<void
     if (!['node', 'jsdom', 'happy-dom'].includes(env))
       throw new Error(`Unsupported environment: ${env}`)
 
-    __vitest_worker__.filepath = file
+    const workerState = getWorkerState()
+    workerState.filepath = file
 
     await withEnv(env as BuiltinEnvironment, config.environmentOptions || {}, async() => {
       await startTests([file], config)
     })
 
-    __vitest_worker__.filepath = undefined
+    workerState.filepath = undefined
   }
 }

packages/vitest/src/node/execute.ts → packages/vitest/src/runtime/execute.ts

@@ -1,15 +1,15 @@
 import { ViteNodeRunner } from 'vite-node/client'
 import type { ModuleCache, ViteNodeRunnerOptions } from 'vite-node'
 import { normalizePath } from 'vite'
-import type { SuiteMocks } from './mocker'
+import type { SuiteMocks } from '../types/mocker'
+import { getWorkerState } from '../utils'
 import { VitestMocker } from './mocker'
 
 export interface ExecuteOptions extends ViteNodeRunnerOptions {
-  files: string[]
   mockMap: SuiteMocks
 }
 
-export async function executeInViteNode(options: ExecuteOptions) {
+export async function executeInViteNode(options: ExecuteOptions & { files: string[] }) {
   const runner = new VitestRunner(options)
 
   // provide the vite define variable in this context
@@ -40,8 +40,10 @@ export class VitestRunner extends ViteNodeRunner {
       this.setCache(dep, module)
     })
 
+    const workerState = getWorkerState()
+
     // support `import.meta.vitest` for test entry
-    if (__vitest_worker__.filepath && normalizePath(__vitest_worker__.filepath) === normalizePath(context.__filename)) {
+    if (workerState.filepath && normalizePath(workerState.filepath) === normalizePath(context.__filename)) {
       // @ts-expect-error injected untyped global
       Object.defineProperty(context.__vite_ssr_import_meta__, 'vitest', { get: () => globalThis.__vitest_index__ })
     }

packages/vitest/src/node/mocker.ts → packages/vitest/src/runtime/mocker.ts

@@ -3,21 +3,13 @@ import { isNodeBuiltin } from 'mlly'
 import { basename, dirname, resolve } from 'pathe'
 import type { ModuleCache } from 'vite-node'
 import { toFilePath } from 'vite-node/utils'
-import { isWindows, mergeSlashes, normalizeId } from '../utils'
+import { getWorkerState, isWindows, mergeSlashes, normalizeId } from '../utils'
 import { distDir } from '../constants'
+import type { PendingSuiteMock } from '../types/mocker'
 import type { ExecuteOptions } from './execute'
 
-export type SuiteMocks = Record<string, Record<string, string | null | (() => unknown)>>
-
 type Callback = (...args: any[]) => unknown
 
-interface PendingSuiteMock {
-  id: string
-  importer: string
-  type: 'mock' | 'unmock'
-  factory?: () => unknown
-}
-
 function getObjectType(value: unknown): string {
   return Object.prototype.toString.apply(value).slice(8, -1)
 }
@@ -40,15 +32,14 @@ function mockPrototype(spyOn: typeof import('../integrations/jest-mock')['spyOn'
   return newProto
 }
 
-const pendingIds: PendingSuiteMock[] = []
-
 export class VitestMocker {
+  private static pendingIds: PendingSuiteMock[] = []
+  private static spyModule?: typeof import('../integrations/jest-mock')
+
   private request!: (dep: string) => unknown
 
   private root: string
-
   private callbacks: Record<string, ((...args: any[]) => unknown)[]> = {}
-  private spy?: typeof import('../integrations/jest-mock')
 
   constructor(
     public options: ExecuteOptions,
@@ -73,7 +64,7 @@ export class VitestMocker {
   }
 
   public getSuiteFilepath(): string {
-    return __vitest_worker__?.filepath || 'global'
+    return getWorkerState().filepath || 'global'
   }
 
   public getMocks() {
@@ -96,15 +87,15 @@ export class VitestMocker {
   }
 
   private async resolveMocks() {
-    await Promise.all(pendingIds.map(async(mock) => {
+    await Promise.all(VitestMocker.pendingIds.map(async(mock) => {
       const { path, external } = await this.resolvePath(mock.id, mock.importer)
       if (mock.type === 'unmock')
         this.unmockPath(path)
       if (mock.type === 'mock')
         this.mockPath(path, external, mock.factory)
     }))
 
-    pendingIds.length = 0
+    VitestMocker.pendingIds = []
   }
 
   private async callFunctionMock(dep: string, mock: () => any) {
@@ -166,7 +157,7 @@ export class VitestMocker {
   }
 
   public mockObject(obj: any) {
-    if (!this.spy)
+    if (!VitestMocker.spyModule)
       throw new Error('Internal Vitest error: Spy function is not defined.')
 
     const type = getObjectType(obj)
@@ -178,7 +169,7 @@ export class VitestMocker {
 
     const newObj = { ...obj }
 
-    const proto = mockPrototype(this.spy.spyOn, Object.getPrototypeOf(obj))
+    const proto = mockPrototype(VitestMocker.spyModule.spyOn, Object.getPrototypeOf(obj))
     Object.setPrototypeOf(newObj, proto)
 
     // eslint-disable-next-line no-restricted-syntax
@@ -187,8 +178,8 @@ export class VitestMocker {
       const type = getObjectType(obj[k])
 
       if (type.includes('Function') && !obj[k]._isMockFunction) {
-        this.spy.spyOn(newObj, k).mockImplementation(() => {})
-        Object.defineProperty(newObj[k], 'length', { value: 0 }) // tinyspy retains length, but jest doesn't
+        VitestMocker.spyModule.spyOn(newObj, k).mockImplementation(() => {})
+        Object.defineProperty(newObj[k], 'length', { value: 0 }) // tinyspy retains length, but jest doesnt
       }
     }
     return newObj
@@ -239,8 +230,8 @@ export class VitestMocker {
   }
 
   private async ensureSpy() {
-    if (this.spy) return
-    this.spy = await this.request(resolve(distDir, 'jest-mock.js')) as typeof import('../integrations/jest-mock')
+    if (VitestMocker.spyModule) return
+    VitestMocker.spyModule = await this.request(resolve(distDir, 'jest-mock.js')) as typeof import('../integrations/jest-mock')
   }
 
   public async requestWithMock(dep: string) {
@@ -268,11 +259,11 @@ export class VitestMocker {
   }
 
   public queueMock(id: string, importer: string, factory?: () => unknown) {
-    pendingIds.push({ type: 'mock', id, importer, factory })
+    VitestMocker.pendingIds.push({ type: 'mock', id, importer, factory })
   }
 
   public queueUnmock(id: string, importer: string) {
-    pendingIds.push({ type: 'unmock', id, importer })
+    VitestMocker.pendingIds.push({ type: 'unmock', id, importer })
   }
 
   public withRequest(request: (dep: string) => unknown) {

packages/vitest/src/runtime/rpc.ts

@@ -1,3 +1,5 @@
+import { getWorkerState } from '../utils'
+
 export const rpc = () => {
-  return __vitest_worker__!.rpc
+  return getWorkerState().rpc
 }

packages/vitest/src/runtime/run.ts

@@ -2,7 +2,7 @@ import { performance } from 'perf_hooks'
 import type { File, HookListener, ResolvedConfig, Suite, SuiteHooks, Task, TaskResult, TaskState, Test } from '../types'
 import { vi } from '../integrations/vi'
 import { getSnapshotClient } from '../integrations/snapshot/chai'
-import { getFullName, hasFailed, hasTests, partitionSuiteChildren } from '../utils'
+import { getFullName, getWorkerState, hasFailed, hasTests, partitionSuiteChildren } from '../utils'
 import { getState, setState } from '../integrations/chai/jest-expect'
 import { takeCoverage } from '../integrations/coverage'
 import { getFn, getHooks } from './map'
@@ -79,7 +79,9 @@ export async function runTest(test: Test) {
 
   getSnapshotClient().setTest(test)
 
-  __vitest_worker__.current = test
+  const workerState = getWorkerState()
+
+  workerState.current = test
 
   try {
     await callSuiteHook(test.suite, test, 'beforeEach', [test, test.suite])
@@ -130,7 +132,7 @@ export async function runTest(test: Test) {
 
   test.result.duration = performance.now() - start
 
-  __vitest_worker__.current = undefined
+  workerState.current = undefined
 
   updateTask(test)
 }
@@ -241,7 +243,7 @@ export async function startTests(paths: string[], config: ResolvedConfig) {
 }
 
 export function clearModuleMocks() {
-  const { clearMocks, mockReset, restoreMocks } = __vitest_worker__.config
+  const { clearMocks, mockReset, restoreMocks } = getWorkerState().config
 
   // since each function calls another, we can just call one
   if (restoreMocks)

packages/vitest/src/runtime/setup.ts

@@ -2,7 +2,7 @@ import { Console } from 'console'
 import { Writable } from 'stream'
 import { environments } from '../integrations/env'
 import type { ResolvedConfig } from '../types'
-import { toArray } from '../utils'
+import { getWorkerState, toArray } from '../utils'
 import * as VitestIndex from '../index'
 import { rpc } from './rpc'
 
@@ -59,7 +59,7 @@ export function setupConsoleLogSpy() {
       rpc().onUserConsoleLog({
         type: 'stdout',
         content: stdoutBuffer.map(i => String(i)).join(''),
-        taskId: __vitest_worker__.current?.id,
+        taskId: getWorkerState().current?.id,
         time: stdoutTime || Date.now(),
       })
     }
@@ -71,7 +71,7 @@ export function setupConsoleLogSpy() {
       rpc().onUserConsoleLog({
         type: 'stderr',
         content: stderrBuffer.map(i => String(i)).join(''),
-        taskId: __vitest_worker__.current?.id,
+        taskId: getWorkerState().current?.id,
         time: stderrTime || Date.now(),
       })
     }
@@ -121,7 +121,7 @@ export async function runSetupFiles(config: ResolvedConfig) {
   const files = toArray(config.setupFiles)
   await Promise.all(
     files.map(async(file) => {
-      __vitest_worker__.moduleCache.delete(file)
+      getWorkerState().moduleCache.delete(file)
       await import(file)
     }),
   )

packages/vitest/src/runtime/worker.ts

@@ -1,15 +1,16 @@
 import { resolve } from 'pathe'
 import { createBirpc } from 'birpc'
-import type { ModuleCache, ResolvedConfig, WorkerContext, WorkerGlobalState, WorkerRPC } from '../types'
+import type { ModuleCache, ResolvedConfig, WorkerContext, WorkerRPC } from '../types'
 import { distDir } from '../constants'
-import { executeInViteNode } from '../node/execute'
+import { getWorkerState } from '../utils'
+import { executeInViteNode } from './execute'
 import { rpc } from './rpc'
 
 let _viteNode: {
   run: (files: string[], config: ResolvedConfig) => Promise<void>
   collect: (files: string[], config: ResolvedConfig) => Promise<void>
 }
-let __vitest_worker__: WorkerGlobalState
+
 const moduleCache: Map<string, ModuleCache> = new Map()
 const mockMap = {}
 
@@ -53,8 +54,9 @@ async function startViteNode(ctx: WorkerContext) {
 }
 
 function init(ctx: WorkerContext) {
-  if (__vitest_worker__ && ctx.config.threads && ctx.config.isolate)
-    throw new Error(`worker for ${ctx.files.join(',')} already initialized by ${__vitest_worker__.ctx.files.join(',')}. This is probably an internal bug of Vitest.`)
+  // @ts-expect-error untyped global
+  if (typeof __vitest_worker__ !== 'undefined' && ctx.config.threads && ctx.config.isolate)
+    throw new Error(`worker for ${ctx.files.join(',')} already initialized by ${getWorkerState().ctx.files.join(',')}. This is probably an internal bug of Vitest.`)
 
   process.stdout.write('\0')
 
@@ -67,6 +69,7 @@ function init(ctx: WorkerContext) {
     ctx,
     moduleCache,
     config,
+    mockMap,
     rpc: createBirpc<WorkerRPC>(
       {},
       {
@@ -93,7 +96,3 @@ export async function run(ctx: WorkerContext) {
   const { run } = await startViteNode(ctx)
   return run(ctx.files, ctx.config)
 }
-
-declare global {
-  let __vitest_worker__: import('vitest').WorkerGlobalState
-}

packages/vitest/src/types/mocker.ts

@@ -0,0 +1,8 @@
+export type SuiteMocks = Record<string, Record<string, string | null | (() => unknown)>>
+
+export interface PendingSuiteMock {
+  id: string
+  importer: string
+  type: 'mock' | 'unmock'
+  factory?: () => unknown
+}

packages/vitest/src/types/worker.ts

@@ -1,6 +1,7 @@
 import type { MessagePort } from 'worker_threads'
 import type { FetchFunction, ModuleCache, RawSourceMap, ViteNodeResolveId } from 'vite-node'
 import type { BirpcReturn } from 'birpc'
+import type { SuiteMocks } from './mocker'
 import type { ResolvedConfig } from './config'
 import type { File, TaskResultPack, Test } from './tasks'
 import type { SnapshotResult } from './snapshot'
@@ -37,4 +38,5 @@ export interface WorkerGlobalState {
   current?: Test
   filepath?: string
   moduleCache: Map<string, ModuleCache>
+  mockMap: SuiteMocks
 }