fix(browser): support module tracking (#8877)

Author: sheremet-va

packages/browser/src/client/public/esm-client-injector.js

@@ -1,26 +1,37 @@
 (() => {
-  const moduleCache = new Map();
-
-  function wrapModule(module) {
-    if (typeof module === "function") {
-      const promise = new Promise((resolve, reject) => {
-        if (typeof __vitest_mocker__ === "undefined")
-          return module().then(resolve, reject);
-        __vitest_mocker__.prepare().finally(() => {
-          module().then(resolve, reject);
-        });
-      });
-      moduleCache.set(promise, { promise, evaluated: false });
-      return promise.finally(() => moduleCache.delete(promise));
+  function wrapModule(moduleCallback) {
+    if (typeof moduleCallback !== "function") {
+      return moduleCallback
+    }
+
+    if (typeof __vitest_mocker__ === "undefined" || typeof __vitest_worker__ === 'undefined') {
+      return moduleCallback()
     }
-    return module;
+
+    const { evaluatedModules } = __vitest_worker__
+    const moduleId = crypto.randomUUID()
+    const viteModule = evaluatedModules.ensureModule(moduleId, moduleId)
+
+    viteModule.evaluated = false
+    viteModule.promise = new Promise((resolve, reject) => {
+      __vitest_mocker__.prepare().finally(() => {
+        moduleCallback().then(resolve, reject)
+      });
+    });
+    return viteModule.promise.finally(() => {
+      viteModule.evaluated = true
+      viteModule.promise = undefined
+
+      evaluatedModules.idToModuleMap.delete(viteModule.id)
+      evaluatedModules.fileToModulesMap.delete(viteModule.file)
+      evaluatedModules.urlToIdModuleMap.delete(viteModule.url)
+    });
   }
 
   window.__vitest_browser_runner__ = {
     wrapModule,
     wrapDynamicImport: wrapModule,
     disposeExceptionTracker: () => {},
-    moduleCache,
     cleanups: [],
     config: { __VITEST_CONFIG__ },
     viteConfig: { __VITEST_VITE_CONFIG__ },

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

@@ -1,5 +1,6 @@
 import type { BrowserRPC } from '@vitest/browser/client'
 import type { WorkerGlobalState } from 'vitest'
+import { EvaluatedModules } from 'vite/module-runner'
 import { getBrowserState } from '../utils'
 
 const config = getBrowserState().config
@@ -30,8 +31,8 @@ const state: WorkerGlobalState = {
     },
   },
   onCleanup: fn => getBrowserState().cleanups.push(fn),
-  evaluatedModules: getBrowserState().evaluatedModules,
-  resolvingModules: getBrowserState().resolvingModules,
+  evaluatedModules: new EvaluatedModules(),
+  resolvingModules: new Set(),
   moduleExecutionInfo: new Map(),
   metaEnv: null as any,
   rpc: null as any,

packages/browser/src/client/utils.ts

@@ -1,5 +1,5 @@
 import type { VitestRunner } from '@vitest/runner'
-import type { EvaluatedModules, SerializedConfig, WorkerGlobalState } from 'vitest'
+import type { SerializedConfig, WorkerGlobalState } from 'vitest'
 import type { IframeOrchestrator } from './orchestrator'
 import type { CommandsManager } from './tester/tester-utils'
 
@@ -65,8 +65,6 @@ export function ensureAwaited<T>(promise: (error?: Error) => Promise<T>): Promis
 export interface BrowserRunnerState {
   files: string[]
   runningFiles: string[]
-  resolvingModules: Set<string>
-  evaluatedModules: EvaluatedModules
   config: SerializedConfig
   provider: string
   runner: VitestRunner

test/browser/src/dynamic.ts

@@ -0,0 +1,3 @@
+window.dispatchEvent(new CustomEvent('module:loaded'))
+
+export {}

test/browser/test/dynamic-module.test.ts

@@ -0,0 +1,9 @@
+import { expect, test, vi } from 'vitest'
+
+test('can await the end of the dynamic import', async () => {
+  const _promise = import('../src/dynamic')
+  const spy = vi.fn()
+  window.addEventListener('module:loaded', spy)
+  await vi.dynamicImportSettled()
+  expect(spy).toHaveBeenCalled()
+})