feat: implement module mocking in browser mode (#5765)

Author: sheremet-va

docs/api/vi.md

@@ -29,10 +29,6 @@ In order to hoist `vi.mock`, Vitest statically analyzes your files. It indicates
 Vitest will not mock modules that were imported inside a [setup file](/config/#setupfiles) because they are cached by the time a test file is running. You can call [`vi.resetModules()`](#vi-resetmodules) inside [`vi.hoisted`](#vi-hoisted) to clear all module caches before running a test file.
 :::
 
-::: warning
-The [browser mode](/guide/browser) does not presently support mocking modules. You can track this feature in the GitHub [issue](https://github.com/vitest-dev/vitest/issues/3046).
-:::
-
 If `factory` is defined, all imports will return its result. Vitest calls factory only once and caches results for all subsequent imports until [`vi.unmock`](#vi-unmock) or [`vi.doUnmock`](#vi-dounmock) is called.
 
 Unlike in `jest`, the factory can be asynchronous. You can use [`vi.importActual`](#vi-importactual) or a helper with the factory passed in as the first argument, and get the original module inside.

docs/config/index.md

@@ -1487,7 +1487,7 @@ Listen to port and serve API. When set to true, the default port is 51204
 
 ### browser {#browser}
 
-- **Type:** `{ enabled?, name?, provider?, headless?, api?, slowHijackESM? }`
+- **Type:** `{ enabled?, name?, provider?, headless?, api? }`
 - **Default:** `{ enabled: false, headless: process.env.CI, api: 63315 }`
 - **CLI:** `--browser`, `--browser=<name>`, `--browser.name=chrome --browser.headless`
 
@@ -1601,17 +1601,6 @@ To have a better type safety when using built-in providers, you can add one of t
 ```
 :::
 
-#### browser.slowHijackESM {#browser-slowhijackesm}
-
-- **Type:** `boolean`
-- **Default:** `false`
-
-When running tests in Node.js Vitest can use its own module resolution to easily mock modules with `vi.mock` syntax. However it's not so easy to replicate ES module resolution in browser, so we need to transform your source files before browser can consume it.
-
-This option has no effect on tests running inside Node.js.
-
-If you rely on spying on ES modules with `vi.spyOn`, you can enable this experimental feature to allow spying on module exports.
-
 #### browser.ui {#browser-ui}
 
 - **Type:** `boolean`

docs/guide/cli-table.md

@@ -55,7 +55,6 @@
 | `--browser.api.strictPort` | Set to true to exit if port is already in use, instead of automatically trying the next available port |
 | `--browser.provider <name>` | Provider used to run browser tests. Some browsers are only available for specific providers. Can be "webdriverio", "playwright", or the path to a custom provider. Visit [`browser.provider`](https://vitest.dev/config/#browser-provider) for more information (default: `"webdriverio"`) |
 | `--browser.providerOptions <options>` | Options that are passed down to a browser provider. Visit [`browser.providerOptions`](https://vitest.dev/config/#browser-provideroptions) for more information |
-| `--browser.slowHijackESM` | Let Vitest use its own module resolution on the browser to enable APIs such as vi.mock and vi.spyOn. Visit [`browser.slowHijackESM`](https://vitest.dev/config/#browser-slowhijackesm) for more information (default: `false`) |
 | `--browser.isolate` | Run every browser test file in isolation. To disable isolation, use `--browser.isolate=false` (default: `true`) |
 | `--pool <pool>` | Specify pool, if not running in the browser (default: `threads`) |
 | `--poolOptions.threads.isolate` | Isolate tests in threads pool (default: `true`) |

packages/browser/context.d.ts

@@ -61,6 +61,10 @@ export const server: {
    * Name of the browser provider.
    */
   provider: string
+  /**
+   * Name of the current browser.
+   */
+  browser: string
   /**
    * Available commands for the browser.
    * @see {@link https://vitest.dev/guide/browser#commands}

packages/browser/package.json

@@ -79,6 +79,8 @@
     "@vitest/ui": "workspace:*",
     "@vitest/ws-client": "workspace:*",
     "@wdio/protocols": "^8.32.0",
+    "birpc": "0.2.17",
+    "flatted": "^3.3.1",
     "periscopic": "^4.0.2",
     "playwright": "^1.44.0",
     "playwright-core": "^1.44.0",

packages/browser/src/client/client.ts

@@ -1,21 +1,116 @@
 import type { CancelReason } from '@vitest/runner'
-import { createClient } from '@vitest/ws-client'
+import { type BirpcReturn, createBirpc } from 'birpc'
+import type { WebSocketBrowserEvents, WebSocketBrowserHandlers } from 'vitest'
+import { parse, stringify } from 'flatted'
+import type { VitestBrowserClientMocker } from './mocker'
+import { getBrowserState } from './utils'
 
 export const PORT = import.meta.hot ? '51204' : location.port
 export const HOST = [location.hostname, PORT].filter(Boolean).join(':')
+export const SESSION_ID = crypto.randomUUID()
 export const ENTRY_URL = `${
   location.protocol === 'https:' ? 'wss:' : 'ws:'
-}//${HOST}/__vitest_api__`
+}//${HOST}/__vitest_browser_api__?type=${getBrowserState().type}&sessionId=${SESSION_ID}`
 
 let setCancel = (_: CancelReason) => {}
 export const onCancel = new Promise<CancelReason>((resolve) => {
   setCancel = resolve
 })
 
-export const client = createClient(ENTRY_URL, {
-  handlers: {
+export interface VitestBrowserClient {
+  rpc: BrowserRPC
+  ws: WebSocket
+  waitForConnection: () => Promise<void>
+}
+
+type BrowserRPC = BirpcReturn<WebSocketBrowserHandlers, WebSocketBrowserEvents>
+
+function createClient() {
+  const autoReconnect = true
+  const reconnectInterval = 2000
+  const reconnectTries = 10
+  const connectTimeout = 60000
+
+  let tries = reconnectTries
+
+  const ctx: VitestBrowserClient = {
+    ws: new WebSocket(ENTRY_URL),
+    waitForConnection,
+  } as VitestBrowserClient
+
+  let onMessage: Function
+
+  ctx.rpc = createBirpc<WebSocketBrowserHandlers, WebSocketBrowserEvents>({
     onCancel: setCancel,
-  },
-})
+    async startMocking(id: string) {
+      // @ts-expect-error not typed global
+      if (typeof __vitest_mocker__ === 'undefined')
+        throw new Error(`Cannot mock modules in the orchestrator process`)
+      // @ts-expect-error not typed global
+      const mocker = __vitest_mocker__ as VitestBrowserClientMocker
+      const exports = await mocker.resolve(id)
+      return Object.keys(exports)
+    },
+  }, {
+    post: msg => ctx.ws.send(msg),
+    on: fn => (onMessage = fn),
+    serialize: e => stringify(e, (_, v) => {
+      if (v instanceof Error) {
+        return {
+          name: v.name,
+          message: v.message,
+          stack: v.stack,
+        }
+      }
+      return v
+    }),
+    deserialize: parse,
+    onTimeoutError(functionName) {
+      throw new Error(`[vitest-browser]: Timeout calling "${functionName}"`)
+    },
+  })
+
+  let openPromise: Promise<void>
+
+  function reconnect(reset = false) {
+    if (reset)
+      tries = reconnectTries
+    ctx.ws = new WebSocket(ENTRY_URL)
+    registerWS()
+  }
+
+  function registerWS() {
+    openPromise = new Promise((resolve, reject) => {
+      const timeout = setTimeout(() => {
+        reject(new Error(`Cannot connect to the server in ${connectTimeout / 1000} seconds`))
+      }, connectTimeout)?.unref?.()
+      if (ctx.ws.OPEN === ctx.ws.readyState)
+        resolve()
+      // still have a listener even if it's already open to update tries
+      ctx.ws.addEventListener('open', () => {
+        tries = reconnectTries
+        resolve()
+        clearTimeout(timeout)
+      })
+    })
+    ctx.ws.addEventListener('message', (v) => {
+      onMessage(v.data)
+    })
+    ctx.ws.addEventListener('close', () => {
+      tries -= 1
+      if (autoReconnect && tries > 0)
+        setTimeout(reconnect, reconnectInterval)
+    })
+  }
+
+  registerWS()
+
+  function waitForConnection() {
+    return openPromise
+  }
+
+  return ctx
+}
 
+export const client = createClient()
 export const channel = new BroadcastChannel('vitest')

packages/browser/src/client/logger.ts

@@ -53,8 +53,11 @@ export async function setupConsoleLogSpy() {
 
   console.trace = (...args: unknown[]) => {
     const content = processLog(args)
-    const error = new Error('Trace')
-    const stack = (error.stack || '').split('\n').slice(2).join('\n')
+    const error = new Error('$$Trace')
+    const stack = (error.stack || '')
+      .split('\n')
+      .slice(error.stack?.includes('$$Trace') ? 2 : 1)
+      .join('\n')
     sendLog('stdout', `${content}\n${stack}`)
     return trace(...args)
   }

packages/browser/src/client/main.ts

@@ -69,12 +69,15 @@ client.ws.addEventListener('open', async () => {
         const filenames = e.data.filenames
         filenames.forEach(filename => runningFiles.delete(filename))
 
-        const iframeId = filenames.length > 1 ? ID_ALL : filenames[0]
-        iframes.get(iframeId)?.remove()
-        iframes.delete(iframeId)
-
-        if (!runningFiles.size)
+        if (!runningFiles.size) {
           await done()
+        }
+        else {
+          // keep the last iframe
+          const iframeId = filenames.length > 1 ? ID_ALL : filenames[0]
+          iframes.get(iframeId)?.remove()
+          iframes.delete(iframeId)
+        }
         break
       }
       // error happened at the top level, this should never happen in user code, but it can trigger during development