feat(vitest): allow using unprocessed environment files (#4098)

Author: sheremet-va

packages/vitest/src/integrations/env/index.ts

@@ -1,6 +1,7 @@
-import { pathToFileURL } from 'node:url'
 import { normalize, resolve } from 'pathe'
 import { resolvePath } from 'mlly'
+import { ViteNodeRunner } from 'vite-node/client'
+import type { ViteNodeRunnerOptions } from 'vite-node'
 import type { BuiltinEnvironment, VitestEnvironment } from '../../types/config'
 import type { Environment } from '../../types'
 import node from './node'
@@ -35,13 +36,26 @@ export function getEnvPackageName(env: VitestEnvironment) {
   return `vitest-environment-${env}`
 }
 
-export async function loadEnvironment(name: VitestEnvironment, root: string): Promise<Environment> {
+const _loaders = new Map<string, ViteNodeRunner>()
+
+export async function createEnvironmentLoader(options: ViteNodeRunnerOptions) {
+  if (!_loaders.has(options.root)) {
+    const loader = new ViteNodeRunner(options)
+    await loader.executeId('/@vite/env')
+    _loaders.set(options.root, loader)
+  }
+  return _loaders.get(options.root)!
+}
+
+export async function loadEnvironment(name: VitestEnvironment, options: ViteNodeRunnerOptions): Promise<Environment> {
   if (isBuiltinEnvironment(name))
     return environments[name]
+  const loader = await createEnvironmentLoader(options)
+  const root = loader.root
   const packageId = name[0] === '.' || name[0] === '/'
     ? resolve(root, name)
     : await resolvePath(`vitest-environment-${name}`, { url: [root] }) ?? resolve(root, name)
-  const pkg = await import(pathToFileURL(normalize(packageId)).href)
+  const pkg = await loader.executeId(normalize(packageId))
   if (!pkg || !pkg.default || typeof pkg.default !== 'object') {
     throw new TypeError(
       `Environment "${name}" is not a valid environment. `

packages/vitest/src/runtime/child.ts

@@ -22,7 +22,7 @@ async function init(ctx: ChildContext) {
     setCancel = resolve
   })
 
-  const rpc = createBirpc<RuntimeRPC, RunnerRPC>(
+  const rpc = createSafeRpc(createBirpc<RuntimeRPC, RunnerRPC>(
     {
       onCancel: setCancel,
     },
@@ -35,9 +35,14 @@ async function init(ctx: ChildContext) {
       },
       on(fn) { process.on('message', fn) },
     },
-  )
+  ))
 
-  const environment = await loadEnvironment(ctx.environment.name, ctx.config.root)
+  const environment = await loadEnvironment(ctx.environment.name, {
+    root: ctx.config.root,
+    fetchModule(id) {
+      return rpc.fetch(id, 'ssr')
+    },
+  })
   if (ctx.environment.transformMode)
     environment.transformMode = ctx.environment.transformMode
 
@@ -52,7 +57,7 @@ async function init(ctx: ChildContext) {
       environment: 0,
       prepare: performance.now(),
     },
-    rpc: createSafeRpc(rpc),
+    rpc,
   }
 
   // @ts-expect-error I know what I am doing :P

packages/vitest/src/runtime/vm.ts

@@ -26,7 +26,7 @@ export async function run(ctx: WorkerContext) {
     setCancel = resolve
   })
 
-  const rpc = createBirpc<RuntimeRPC>(
+  const rpc = createSafeRpc(createBirpc<RuntimeRPC>(
     {
       onCancel: setCancel,
     },
@@ -35,9 +35,14 @@ export async function run(ctx: WorkerContext) {
       post(v) { port.postMessage(v) },
       on(fn) { port.addListener('message', fn) },
     },
-  )
+  ))
 
-  const environment = await loadEnvironment(ctx.environment.name, ctx.config.root)
+  const environment = await loadEnvironment(ctx.environment.name, {
+    root: ctx.config.root,
+    fetchModule(id) {
+      return rpc.fetch(id, 'ssr')
+    },
+  })
 
   if (!environment.setupVM) {
     const envName = ctx.environment.name
@@ -59,7 +64,7 @@ export async function run(ctx: WorkerContext) {
       environment: performance.now(),
       prepare: performance.now(),
     },
-    rpc: createSafeRpc(rpc),
+    rpc,
   }
 
   installSourcemapsSupport({

packages/vitest/src/runtime/worker.ts

@@ -24,7 +24,7 @@ async function init(ctx: WorkerContext) {
     setCancel = resolve
   })
 
-  const rpc = createBirpc<RuntimeRPC, RunnerRPC>(
+  const rpc = createSafeRpc(createBirpc<RuntimeRPC, RunnerRPC>(
     {
       onCancel: setCancel,
     },
@@ -33,9 +33,14 @@ async function init(ctx: WorkerContext) {
       post(v) { port.postMessage(v) },
       on(fn) { port.addListener('message', fn) },
     },
-  )
+  ))
 
-  const environment = await loadEnvironment(ctx.environment.name, ctx.config.root)
+  const environment = await loadEnvironment(ctx.environment.name, {
+    root: ctx.config.root,
+    fetchModule(id) {
+      return rpc.fetch(id, 'ssr')
+    },
+  })
   if (ctx.environment.transformMode)
     environment.transformMode = ctx.environment.transformMode
 
@@ -50,7 +55,7 @@ async function init(ctx: WorkerContext) {
       environment: 0,
       prepare: performance.now(),
     },
-    rpc: createSafeRpc(rpc),
+    rpc,
   }
 
   // @ts-expect-error I know what I am doing :P

test/env-custom/vitest-environment-custom/index.mjs renamed to test/env-custom/vitest-environment-custom/index.ts

@@ -1,6 +1,7 @@
 import vm from 'node:vm'
+import type { Environment } from 'vitest'
 
-export default {
+export default <Environment>{
   name: 'custom',
   transformMode: 'ssr',
   setupVM({ custom }) {

test/env-custom/vitest-environment-custom/package.json

@@ -2,6 +2,7 @@
   "name": "vitest-environment-custom",
   "private": true,
   "exports": {
-    ".": "./index.mjs"
-  }
+    ".": "./index.ts"
+  },
+  "main": "index.ts"
 }