feat: support autospy in the browser

Author: sheremet-va

packages/browser/src/client/channel.ts

@@ -26,6 +26,7 @@ export interface IframeMockEvent {
   type: 'mock'
   paths: string[]
   mock: string | undefined | null
+  behaviour: 'autospy' | 'automock' | 'manual'
 }
 
 export interface IframeUnmockEvent {

packages/browser/src/client/tester/mocker.ts

@@ -14,6 +14,7 @@ interface SpyModule {
 export class VitestBrowserClientMocker {
   private queue = new Set<Promise<void>>()
   private mocks: Record<string, undefined | null | string> = {}
+  private behaviours: Record<string, 'autospy' | 'automock' | 'manual'> = {}
   private mockObjects: Record<string, any> = {}
   private factories: Record<string, () => any> = {}
   private ids = new Set<string>()
@@ -91,14 +92,20 @@ export class VitestBrowserClientMocker {
       return await this.resolve(resolvedId)
     }
 
-    if (type === 'redirect') {
-      const url = new URL(`/@id/${mockPath}`, location.href)
-      return import(/* @vite-ignore */ url.toString())
+    const behavior = this.behaviours[resolvedId] || 'automock'
+
+    if (type === 'automock' || behavior === 'autospy') {
+      const url = new URL(`/@id/${resolvedId}`, location.href)
+      const query = url.search ? `${url.search}&t=${now()}` : `?t=${now()}`
+      const moduleObject = await import(/* @vite-ignore */ `${url.pathname}${query}&mock=${behavior}${url.hash}`)
+      return this.mockObject(moduleObject, {}, behavior)
+    }
+
+    if (typeof mockPath !== 'string') {
+      throw new TypeError(`Mock path is not a string: ${mockPath}. This is a bug in Vitest. Please, open a new issue with reproduction.`)
     }
-    const url = new URL(`/@id/${resolvedId}`, location.href)
-    const query = url.search ? `${url.search}&t=${now()}` : `?t=${now()}`
-    const moduleObject = await import(/* @vite-ignore */ `${url.pathname}${query}${url.hash}`)
-    return this.mockObject(moduleObject)
+    const url = new URL(`/@id/${mockPath}`, location.href)
+    return import(/* @vite-ignore */ url.toString())
   }
 
   public getMockContext() {
@@ -142,32 +149,35 @@ export class VitestBrowserClientMocker {
     }
   }
 
-  public queueMock(id: string, importer: string, factory?: () => any) {
+  public queueMock(id: string, importer: string, factoryOrOptions?: MockOptions | (() => any)) {
     const promise = rpc()
-      .resolveMock(id, importer, !!factory)
+      .resolveMock(id, importer, typeof factoryOrOptions === 'function')
       .then(async ({ mockPath, resolvedId, needsInterop }) => {
         this.ids.add(resolvedId)
         const urlPaths = resolveMockPaths(cleanVersion(resolvedId))
         const resolvedMock
           = typeof mockPath === 'string'
             ? new URL(resolvedMockedPath(cleanVersion(mockPath)), location.href).toString()
             : mockPath
-        const _factory = factory && needsInterop
+        const factory = typeof factoryOrOptions === 'function'
           ? async () => {
-            const data = await factory()
-            return { default: data }
+            const data = await factoryOrOptions()
+            return needsInterop ? { default: data } : data
           }
-          : factory
+          : undefined
+        const behaviour = getMockBehaviour(factoryOrOptions)
         urlPaths.forEach((url) => {
           this.mocks[url] = resolvedMock
-          if (_factory) {
-            this.factories[url] = _factory
+          if (factory) {
+            this.factories[url] = factory
           }
+          this.behaviours[url] = behaviour
         })
         channel.postMessage({
           type: 'mock',
           paths: urlPaths,
           mock: resolvedMock,
+          behaviour,
         })
         await waitForChannel('mock:done')
       })
@@ -214,6 +224,7 @@ export class VitestBrowserClientMocker {
   public mockObject(
     object: Record<Key, any>,
     mockExports: Record<Key, any> = {},
+    behaviour: 'autospy' | 'automock' | 'manual' = 'automock',
   ) {
     const finalizers = new Array<() => void>()
     const refs = new RefTracker()
@@ -330,13 +341,14 @@ export class VitestBrowserClientMocker {
               }
             }
           }
-          const mock = spyModule
-            .spyOn(newContainer, property)
-            .mockImplementation(mockFunction)
-          mock.mockRestore = () => {
-            mock.mockReset()
+          const mock = spyModule.spyOn(newContainer, property)
+          if (behaviour === 'automock') {
             mock.mockImplementation(mockFunction)
-            return mock
+            mock.mockRestore = () => {
+              mock.mockReset()
+              mock.mockImplementation(mockFunction)
+              return mock
+            }
           }
           // tinyspy retains length, but jest doesn't.
           Object.defineProperty(newContainer[property], 'length', { value: 0 })
@@ -465,3 +477,19 @@ const versionRegexp = /(\?|&)v=\w{8}/
 function cleanVersion(url: string) {
   return url.replace(versionRegexp, '')
 }
+
+export interface MockOptions {
+  spy?: boolean
+}
+
+export type MockBehaviour = 'autospy' | 'automock' | 'manual'
+
+function getMockBehaviour(factoryOrOptions?: (() => void) | MockOptions): MockBehaviour {
+  if (!factoryOrOptions) {
+    return 'automock'
+  }
+  if (typeof factoryOrOptions === 'function') {
+    return 'manual'
+  }
+  return factoryOrOptions.spy ? 'autospy' : 'automock'
+}

packages/browser/src/client/tester/msw.ts

@@ -7,7 +7,10 @@ import type {
 } from '@vitest/browser/client'
 
 export function createModuleMocker() {
-  const mocks: Map<string, string | null | undefined> = new Map()
+  const mocks: Map<string, {
+    mock: string | null | undefined
+    behaviour: 'automock' | 'autospy' | 'manual'
+  }> = new Map()
 
   let started = false
   let startPromise: undefined | Promise<unknown>
@@ -34,7 +37,7 @@ export function createModuleMocker() {
             return passthrough()
           }
 
-          const mock = mocks.get(path)
+          const { mock, behaviour } = mocks.get(path)!
 
           // using a factory
           if (mock === undefined) {
@@ -56,11 +59,11 @@ export function createModuleMocker() {
             })
           }
 
-          if (typeof mock === 'string') {
-            return Response.redirect(mock)
+          if (behaviour === 'autospy' || mock === null) {
+            return Response.redirect(injectQuery(path, `mock=${behaviour}`))
           }
 
-          return Response.redirect(injectQuery(path, 'mock=auto'))
+          return Response.redirect(mock)
         }),
       )
       return worker.start({
@@ -80,7 +83,7 @@ export function createModuleMocker() {
   return {
     async mock(event: IframeMockEvent) {
       await init()
-      event.paths.forEach(path => mocks.set(path, event.mock))
+      event.paths.forEach(path => mocks.set(path, { mock: event.mock, behaviour: event.behaviour }))
       channel.postMessage(<IframeMockingDoneEvent>{ type: 'mock:done' })
     },
     async unmock(event: IframeUnmockEvent) {

packages/vitest/src/node/automock.ts

@@ -12,7 +12,7 @@ import type {
 import MagicString from 'magic-string'
 
 // TODO: better source map replacement
-export function automockModule(code: string, parse: (code: string) => Program) {
+export function automockModule(code: string, behavior: 'automock' | 'autospy', parse: (code: string) => Program) {
   const ast = parse(code)
 
   const m = new MagicString(code)
@@ -145,7 +145,7 @@ const __vitest_es_current_module__ = {
   __esModule: true,
   ${allSpecifiers.map(({ name }) => `["${name}"]: ${name},`).join('\n  ')}
 }
-const __vitest_mocked_module__ = __vitest_mocker__.mockObject(__vitest_es_current_module__)
+const __vitest_mocked_module__ = __vitest_mocker__.mockObject(__vitest_es_current_module__, {}, "${behavior}")
 `
   const assigning = allSpecifiers
     .map(({ name }, index) => {

packages/vitest/src/node/plugins/mocks.ts

@@ -20,8 +20,9 @@ export function MocksPlugins(): Plugin[] {
       name: 'vitest:automock',
       enforce: 'post',
       transform(code, id) {
-        if (id.includes('mock=auto')) {
-          const ms = automockModule(code, this.parse)
+        if (id.includes('mock=automock') || id.includes('mock=autospy')) {
+          const behavior = id.includes('mock=automock') ? 'automock' : 'autospy'
+          const ms = automockModule(code, behavior, this.parse)
           return {
             code: ms.toString(),
             map: ms.generateMap({ hires: true, source: cleanUrl(id) }),

test/browser/fixtures/mocking/autospying.test.ts

@@ -0,0 +1,16 @@
+import { expect, test, vi } from 'vitest'
+import { calculator } from './src/calculator'
+import * as mocks_calculator from './src/mocks_calculator'
+
+vi.mock('./src/calculator', { spy: true })
+vi.mock('./src/mocks_calculator', { spy: true })
+
+test('correctly spies on a regular module', () => {
+  expect(calculator('plus', 1, 2)).toBe(3)
+  expect(calculator).toHaveBeenCalled()
+})
+
+test('spy options overrides __mocks__ folder', () => {
+  expect(mocks_calculator.calculator('plus', 1, 2)).toBe(3)
+  expect(mocks_calculator.calculator).toHaveBeenCalled()
+})