fix: add proxyHijackESM for better spyOn in browser

examples/react/test/basic.test.tsx

@@ -1,6 +1,7 @@
 import React from 'react'
 import renderer from 'react-test-renderer'
 import Link from '../components/Link.jsx'
+import * as linkModule from '../components/Link.jsx'
 
 function toJson(component: renderer.ReactTestRenderer) {
   const result = component.toJSON()
@@ -29,3 +30,19 @@ test('Link changes the class when hovered', () => {
   tree = toJson(component)
   expect(tree).toMatchSnapshot()
 })
+
+test('Link can be spied', () => {
+  vi.spyOn(linkModule, 'default').mockImplementation(() => {
+    return <div>Hello</div>
+  })
+
+  const component = renderer.create(
+    <Link page="http://antfu.me">Anthony Fu</Link>,
+  )
+
+  const tree = toJson(component)
+  console.warn('tree', tree)
+
+  expect(tree.type).toBe('div')
+  expect(tree.children).toStrictEqual(['Hello'])
+})

examples/react/vitest.config.ts

@@ -6,5 +6,14 @@ export default defineConfig({
   test: {
     globals: true,
     environment: 'happy-dom',
+    browser: {
+      enabled: true,
+      name: process.env.BROWSER || 'chrome',
+      headless: false,
+      provider: process.env.PROVIDER || 'webdriverio',
+      isolate: false,
+      proxyHijackESM: true,
+    },
+
   },
 })

packages/browser/src/client/fakeModule.ts

@@ -0,0 +1,62 @@
+// this is just here so Proxy knows we're wrapping a function
+function neverFunctionForProxy() {
+  throw new Error(`should not get here`)
+}
+
+// allows replacing 'function' types on this module: passes through to the 'current value'
+function prepareFakeModuleForSpy(fakeModule: any, contents: any) {
+  for (const key in contents) {
+    if (typeof contents[key] !== 'function') {
+      fakeModule[key] = contents[key]
+      continue
+    }
+
+    // uses Reflect to implement Proxy, but use 'current value' and not the original target
+    // this is a Proxy pointing _to_ Reflect, which means it "supports all operations" but on the dynamic target
+    const currentValueReflect = new Proxy(Reflect, {
+      get(_reflect, reflectKey) {
+        return (...args: any[]) => (Reflect as any)[reflectKey](contents[key], ...args.slice(1))
+      },
+    })
+    fakeModule[key] = new Proxy(neverFunctionForProxy, currentValueReflect)
+  }
+}
+
+// this enables bound-like named imports with the import rewriting scheme
+export function buildFakeModule<T extends Record<string | symbol, any>>(module: T): T {
+  const fakeModule: Record<string | symbol, any> = { [Symbol.toStringTag]: 'Module' }
+  const contents: Record<string | symbol, any> = { ...module }
+
+  prepareFakeModuleForSpy(fakeModule, contents)
+
+  // this intercepts tinyspy which uses Object.defineProperty to rewrite an object (even a "module")
+  return new Proxy({}, {
+    ownKeys(_target) {
+      return Reflect.ownKeys(fakeModule)
+    },
+
+    set(_target, _property) {
+      return false
+    },
+
+    get(_target, property) {
+      return fakeModule[property]
+    },
+
+    getOwnPropertyDescriptor(_target, p) {
+      return Reflect.getOwnPropertyDescriptor(contents, p)
+    },
+
+    defineProperty(_target, property, attributes) {
+      if (attributes.get || attributes.set)
+        throw new Error(`can't defineProperty with get/set on fake module`)
+
+      if ('value' in attributes) {
+        contents[property] = attributes.value
+        return true
+      }
+
+      return false
+    },
+  }) as T
+}

packages/browser/src/client/index.html

@@ -63,6 +63,11 @@
         }
       }
 
+      // blindly load all code before mocker arrives
+      window.__vitest_mocker__ = {
+        wrap: (fn) => fn(),
+      }
+
       window.__vi_export_all__ = exportAll
 
       // TODO: allow easier rewriting of import.meta.env

packages/browser/src/client/main.ts

@@ -18,6 +18,7 @@ export const ENTRY_URL = `${
 
 let config: ResolvedConfig | undefined
 let runner: VitestRunner | undefined
+let mocker: VitestBrowserClientMocker
 const browserHashMap = new Map<string, [test: boolean, timestamp: string]>()
 
 const url = new URL(location.href)
@@ -173,8 +174,10 @@ ws.addEventListener('open', async () => {
     },
     providedContext: await client.rpc.getProvidedContext(),
   }
+
+  mocker = new VitestBrowserClientMocker(config!)
   // @ts-expect-error mocking vitest apis
-  globalThis.__vitest_mocker__ = new VitestBrowserClientMocker()
+  globalThis.__vitest_mocker__ = mocker
 
   const paths = getQueryPaths()
 
@@ -224,6 +227,7 @@ async function runTests(paths: string[], config: ResolvedConfig) {
     preparedData = await prepareTestEnvironment(config)
   }
   catch (err) {
+    console.warn('Could not prepareTestEnvironment', err)
     location.reload()
     return
   }
@@ -250,8 +254,14 @@ async function runTests(paths: string[], config: ResolvedConfig) {
 
     runningTests = true
 
-    for (const file of files)
-      await startTests([file], runner)
+    for (const file of files) {
+      try {
+        await startTests([file], runner)
+      }
+      finally {
+        mocker.resetAfterFile()
+      }
+    }
   }
   finally {
     runningTests = false

packages/browser/src/client/mocker.ts

@@ -1,8 +1,40 @@
+import type { ResolvedConfig } from 'vitest'
+import { buildFakeModule } from './fakeModule'
+
 function throwNotImplemented(name: string) {
   throw new Error(`[vitest] ${name} is not implemented in browser environment yet.`)
 }
 
 export class VitestBrowserClientMocker {
+  constructor(public config: ResolvedConfig) {}
+
+  private wrappedImports = new WeakMap<any, any>()
+
+  /**
+   * Browser tests don't run in parallel. This clears all mocks after each run.
+   */
+  public resetAfterFile() {
+    this.resetModules()
+  }
+
+  public resetModules() {
+    this.wrappedImports = new WeakMap()
+  }
+
+  public async wrap(fn: () => Promise<any>) {
+    if (!this.config.browser.proxyHijackESM)
+      throw new Error(`hijackESM disabled but mocker invoked`)
+
+    const module = await fn()
+
+    let wrapped = this.wrappedImports.get(module)
+    if (wrapped === undefined) {
+      wrapped = buildFakeModule(module)
+      this.wrappedImports.set(module, wrapped)
+    }
+    return wrapped
+  }
+
   public importActual() {
     throwNotImplemented('importActual')
   }

packages/browser/src/node/esmProxy.ts

@@ -0,0 +1,54 @@
+import MagicString from 'magic-string'
+import type { PluginContext } from 'rollup'
+import type { Expression } from 'estree'
+import type { Positioned } from './esmWalker'
+import { esmWalker } from './esmWalker'
+
+// don't allow mocking vitest itself
+// (this *also* changes the position of user code in the getImporter() helper)
+const skipImports = [
+  '^vitest$',
+  '^@vitest/',
+]
+
+export async function insertEsmProxy(
+  code: string,
+  id: string,
+  parse: PluginContext['parse'],
+) {
+  const s = new MagicString(code)
+
+  let ast: any
+  try {
+    ast = parse(code)
+  }
+  catch (err) {
+    console.error(`Cannot parse ${id}:\n${(err as any).message}`)
+    return
+  }
+
+  esmWalker(ast, {
+    onDynamicImport(node) {
+      const expression = (node.source as Positioned<Expression>)
+      if (!(expression.type === 'Literal' && typeof expression.value === 'string'))
+        // this is a non-string import so vite won't change it anyway
+        return
+
+      const value = expression.value as string
+
+      if (skipImports.some(i => value.match(i)))
+        return
+
+      const replace = `__vitest_mocker__.wrap(() => import(${JSON.stringify(value)}))`
+      s.overwrite(node.start, node.end, replace)
+    },
+    onIdentifier() {},
+    onImportMeta() {},
+  })
+
+  return {
+    ast,
+    code: s.toString(),
+    map: s.generateMap({ hires: 'boundary', source: id }),
+  }
+}

packages/browser/src/node/index.ts

@@ -6,8 +6,9 @@ import sirv from 'sirv'
 import type { Plugin } from 'vite'
 import type { WorkspaceProject } from 'vitest/node'
 import { injectVitestModule } from './esmInjector'
+import { insertEsmProxy } from './esmProxy'
 
-export default (project: WorkspaceProject, base = '/'): Plugin[] => {
+export default (existingPlugins: Plugin[], project: WorkspaceProject, base = '/'): Plugin[] => {
   const pkgRoot = resolve(fileURLToPath(import.meta.url), '../..')
   const distRoot = resolve(pkgRoot, 'dist')
 
@@ -98,5 +99,19 @@ export default (project: WorkspaceProject, base = '/'): Plugin[] => {
         return injectVitestModule(source, id, this.parse)
       },
     },
+
+    // MocksPlugin needs to be placed here for proxyHijackESM to work
+    ...existingPlugins,
+
+    {
+      name: 'vitest:browser:esm-proxy',
+      enforce: 'post',
+      async transform(source, id) {
+        const proxyHijackESM = project.config.browser.proxyHijackESM ?? false
+        if (!proxyHijackESM)
+          return
+        return insertEsmProxy(source, id, this.parse)
+      },
+    },
   ]
 }

packages/vitest/src/integrations/browser/server.ts

@@ -1,4 +1,5 @@
 import { createServer } from 'vite'
+import type { Plugin } from 'vite'
 import { defaultBrowserPort } from '../../constants'
 import { ensurePackageInstalled } from '../../node/pkg'
 import { resolveApiServerConfig } from '../../node/config'
@@ -8,11 +9,54 @@ import { MocksPlugin } from '../../node/plugins/mocks'
 import { resolveFsAllow } from '../../node/plugins/utils'
 
 export async function createBrowserServer(project: WorkspaceProject, configFile: string | undefined) {
+  if (project.config.browser.proxyHijackESM && project.config.browser.slowHijackESM)
+    throw new Error(`cannot set both proxyHijackESM and slowHijackESM`)
+
   const root = project.config.root
 
   await ensurePackageInstalled('@vitest/browser', root)
 
   const configPath = typeof configFile === 'string' ? configFile : false
+  const alwaysMock = Boolean(project.config.browser.proxyHijackESM)
+
+  const existingPlugins: Plugin[] = [
+    CoverageTransform(project.ctx),
+    {
+      enforce: 'post',
+      name: 'vitest:browser:config',
+      async config(config) {
+        const server = resolveApiServerConfig(config.test?.browser || {}) || {
+          port: defaultBrowserPort,
+        }
+
+        // browser never runs in middleware mode
+        server.middlewareMode = false
+
+        config.server = {
+          ...config.server,
+          ...server,
+        }
+        config.server.fs ??= {}
+        config.server.fs.allow = config.server.fs.allow || []
+        config.server.fs.allow.push(
+          ...resolveFsAllow(
+            project.ctx.config.root,
+            project.ctx.server.config.configFile,
+          ),
+        )
+
+        return {
+          resolve: {
+            alias: config.test?.alias,
+          },
+          server: {
+            watch: null,
+          },
+        }
+      },
+    },
+    MocksPlugin({ always: alwaysMock }),
+  ]
 
   const server = await createServer({
     logLevel: 'error',
@@ -25,45 +69,7 @@ export async function createBrowserServer(project: WorkspaceProject, configFile:
         ignored: ['**/**'],
       },
     },
-    plugins: [
-      (await import('@vitest/browser')).default(project, '/'),
-      CoverageTransform(project.ctx),
-      {
-        enforce: 'post',
-        name: 'vitest:browser:config',
-        async config(config) {
-          const server = resolveApiServerConfig(config.test?.browser || {}) || {
-            port: defaultBrowserPort,
-          }
-
-          // browser never runs in middleware mode
-          server.middlewareMode = false
-
-          config.server = {
-            ...config.server,
-            ...server,
-          }
-          config.server.fs ??= {}
-          config.server.fs.allow = config.server.fs.allow || []
-          config.server.fs.allow.push(
-            ...resolveFsAllow(
-              project.ctx.config.root,
-              project.ctx.server.config.configFile,
-            ),
-          )
-
-          return {
-            resolve: {
-              alias: config.test?.alias,
-            },
-            server: {
-              watch: null,
-            },
-          }
-        },
-      },
-      MocksPlugin(),
-    ],
+    plugins: (await import('@vitest/browser')).default(existingPlugins, project, '/'),
   })
 
   await server.listen()