Fixes for module mocking (#9853)

* allow mocking of prebundled deps

* improved TS support for mock primitives

* export all mock types

* don't run getTemplate in windows

* reduce coverage a bit

packages/wdio-browser-runner/src/browser/spy.ts

@@ -1,13 +1,14 @@
+import type { MaybeMocked } from '@vitest/spy'
+
 import { MESSAGE_TYPES } from '../constants.js'
-import type { SocketMessage, SocketMessagePayload } from '../vite/types.js'
+import type { MockFactoryWithHelper } from '../types'
+import type { SocketMessage, SocketMessagePayload } from '../vite/types'
 
 /**
  * re-export mock module
  */
 export * from '@vitest/spy'
 
-type MockFactoryWithHelper = (importOriginal: <T = unknown>() => Promise<T>) => any;
-
 const a = document.createElement('a')
 function resolveUrl(path: string) {
     a.href = path
@@ -66,3 +67,8 @@ socket.addEventListener('message', (ev) => {
         // ignore
     }
 })
+
+/**
+ * utility helper for type conversions
+ */
+export function mocked<T>(item: T) { return item as any as MaybeMocked<T> }

packages/wdio-browser-runner/src/index.ts

@@ -2,7 +2,6 @@ import fs from 'node:fs/promises'
 import path from 'node:path'
 
 import logger from '@wdio/logger'
-import type { RunArgs, WorkerInstance } from '@wdio/local-runner'
 import LocalRunner from '@wdio/local-runner'
 import { attach } from 'webdriverio'
 import libCoverage, { type CoverageMap, type CoverageMapData } from 'istanbul-lib-coverage'
@@ -10,8 +9,10 @@ import libReport from 'istanbul-lib-report'
 import libSourceMap from 'istanbul-lib-source-maps'
 import reports from 'istanbul-reports'
 
+import type { RunArgs, WorkerInstance } from '@wdio/local-runner'
 import type { SessionStartedMessage, SessionEndedMessage, WorkerHookResultMessage, WorkerCoverageMapMessage } from '@wdio/runner'
 import type { Options } from '@wdio/types'
+import type { MaybeMocked, MaybeMockedDeep, MaybePartiallyMocked, MaybePartiallyMockedDeep } from '@vitest/spy'
 
 import { ViteServer } from './vite/server.js'
 import {
@@ -20,7 +21,7 @@ import {
 } from './constants.js'
 import { makeHeadless, getCoverageByFactor } from './utils.js'
 import type { HookTriggerEvent } from './vite/types.js'
-import type { BrowserRunnerOptions as BrowserRunnerOptionsImport, CoverageOptions } from './types.js'
+import type { BrowserRunnerOptions as BrowserRunnerOptionsImport, CoverageOptions, MockFactoryWithHelper } from './types.js'
 
 const log = logger('@wdio/browser-runner')
 export default class BrowserRunner extends LocalRunner {
@@ -213,3 +214,70 @@ declare global {
  * re-export mock types
  */
 export * from '@vitest/spy'
+
+/**
+ * The following exports are meaningless and only there to allow proper type support.
+ * The actual implementation can be found in /src/browser.spy.ts
+ */
+
+/**
+ * Makes all imports to passed module to be mocked.
+ *
+ * If there is a factory, will return it's result. The call to `mock` is hoisted to the top of the file,
+ * so you don't have access to variables declared in the global file scope, if you didn't put them before imports!
+ *
+ * If __mocks__ folder with file of the same name exist, all imports will return it.
+ *
+ * If there is no __mocks__ folder or a file with the same name inside, will call original module and mock it.
+ *
+ * @param {string} path Path to the module.
+ * @param {MockFactoryWithHelper} factory (optional) Factory for the mocked module. Has the highest priority.
+ */
+// eslint-disable-next-line @typescript-eslint/no-unused-vars
+export function mock (path: string, factory?: MockFactoryWithHelper) {}
+
+/**
+ * Removes module from mocked registry. All subsequent calls to import will return original module even if it was mocked.
+ *
+ * @param path Path to the module.
+ */
+// eslint-disable-next-line @typescript-eslint/no-unused-vars
+export function unmock(moduleName: string) {}
+
+/**
+ * Type helpers for TypeScript. In reality just returns the object that was passed.
+ * @example
+ * import example from './example'
+ * vi.mock('./example')
+ *
+ * test('1+1 equals 2' async () => {
+ *  vi.mocked(example.calc).mockRestore()
+ *
+ *  const res = example.calc(1, '+', 1)
+ *
+ *  expect(res).toBe(2)
+ * })
+ *
+ * @param item Anything that can be mocked
+ * @returns
+ */
+export function mocked<T>(item: T, deep?: true): MaybeMockedDeep<T>
+export function mocked<T>(item: T, deep?: false): MaybeMockedDeep<T>
+export function mocked<T>(item: T, options: {
+    partial?: false;
+    deep?: false;
+}): MaybeMocked<T>;
+export function mocked<T>(item: T, options: {
+    partial?: false;
+    deep: true;
+}): MaybeMockedDeep<T>;
+export function mocked<T>(item: T, options: {
+    partial: true;
+    deep?: false;
+}): MaybePartiallyMocked<T>;
+export function mocked<T>(item: T, options: {
+    partial: true;
+    deep: true;
+}): MaybePartiallyMockedDeep<T>
+// eslint-disable-next-line @typescript-eslint/no-unused-vars
+export function mocked (item: any, options: any) {}

packages/wdio-browser-runner/src/types.ts

@@ -120,3 +120,5 @@ export interface Environment {
     sessionId: string
     injectGlobals: boolean
 }
+
+export type MockFactoryWithHelper = (importOriginal: <T = unknown>() => Promise<T>) => any;

packages/wdio-browser-runner/src/vite/plugins/mockHoisting.ts

@@ -32,11 +32,14 @@ export function mockHoisting(mockHandler: MockHandler): Plugin[] {
                 }
             }
 
+            const preBundledDepName = path.basename(id).split('?')[0]
             const mockedMod = (
                 // mocked file
                 mockHandler.mocks.get(os.platform() === 'win32' ? `/${id}` : id) ||
                 // mocked dependency
-                mockHandler.mocks.get(path.basename(id, path.extname(id)))
+                mockHandler.mocks.get(path.basename(id, path.extname(id))) ||
+                // pre-bundled deps e.g. /node_modules/.vite/deps/algoliasearch_lite.js?v=e31c24e
+                [...mockHandler.mocks.values()].find((mock) => `${mock.path.replace('/', '_')}.js` === preBundledDepName)
             )
             if (mockedMod) {
                 const newCode = mockedMod.namedExports.map((ne) => {

packages/wdio-browser-runner/tests/vite/plugins/__snapshots__/mockHoisting.test.ts.snap

@@ -6,6 +6,8 @@ export const bar = window.__wdioMockFactories__['mocked']['bar'];
 export default window.__wdioMockFactories__['mocked'].default;"
 `;
 
+exports[`returns mock for prebundled dep 1`] = `"export default window.__wdioMockFactories__['algoliasearch/lite'].default;"`;
+
 exports[`transforms test file properly for mocking 1`] = `
 PrintResult {
   "code": "import { spyOn, mock, unmock as foobar, fn } from '@wdio/browser-runner'

packages/wdio-browser-runner/tests/vite/plugins/mockHoisting.test.ts

@@ -15,10 +15,17 @@ const mockHandler: any = {
     resetMocks: vi.fn(),
     manualMocks: [],
     unmock: vi.fn(),
-    mocks: { get: vi.fn().mockReturnValue({
-        namedExports: ['foo', 'bar', 'default'],
-        path: 'mocked'
-    }) }
+    mocks: {
+        get: vi.fn().mockReturnValue({
+            namedExports: ['foo', 'bar', 'default'],
+            path: 'mocked'
+        }),
+        values: vi.fn().mockReturnValue([{
+            path: 'algoliasearch/lite',
+            origin: '/some/path',
+            namedExports: ['default']
+        }])
+    }
 }
 
 test('exposes correct format', () => {
@@ -68,3 +75,10 @@ test('returns mock file', async () => {
     const prePlugin = mockHoisting(mockHandler).shift()!
     expect(await (prePlugin.load as Function)('foobar')).toMatchSnapshot()
 })
+
+test('returns mock for prebundled dep', async () => {
+    vi.mocked(mockHandler.mocks.get).mockReturnValue()
+    const id = '/path/to/project/node_modules/.vite/deps/algoliasearch_lite.js?v=e31c24e7'
+    const prePlugin = mockHoisting(mockHandler).shift()!
+    expect(await (prePlugin.load as Function)(id)).toMatchSnapshot()
+})

packages/wdio-browser-runner/tests/vite/utils.test.ts

@@ -16,27 +16,22 @@ vi.mock('fakeDep', () => ({
     default: 'I am fake'
 }))
 
-describe('getTemplate', () => {
-    it('fails if vue helpers are not installed', async () => {
-        // skip for Windows
-        if (os.platform() === 'win32') {
-            return
-        }
-        vi.mocked(resolve).mockRejectedValue(new Error('not there'))
-        await expect(getTemplate({ preset: 'vue' }, {} as any, ''))
-            .rejects.toThrow(/Fail to set-up Vue environment/)
-    })
+// skip for Windows
+if (os.platform() !== 'win32') {
+    describe('getTemplate', () => {
+        it('fails if vue helpers are not installed', async () => {
+            vi.mocked(resolve).mockRejectedValue(new Error('not there'))
+            await expect(getTemplate({ preset: 'vue' }, {} as any, ''))
+                .rejects.toThrow(/Fail to set-up Vue environment/)
+        })
 
-    it('renders template correctly', async () => {
-        // skip for Windows
-        if (os.platform() === 'win32') {
-            return
-        }
-        vi.mocked(resolve).mockResolvedValue('file:///foo/bar/vue')
-        expect(await getTemplate({ preset: 'vue' }, {} as any, '/spec.js', { some: 'env' })).toMatchSnapshot()
-        expect(fs.readFile).toBeCalledTimes(2)
+        it('renders template correctly', async () => {
+            vi.mocked(resolve).mockResolvedValue('file:///foo/bar/vue')
+            expect(await getTemplate({ preset: 'vue' }, {} as any, '/spec.js', { some: 'env' })).toMatchSnapshot()
+            expect(fs.readFile).toBeCalledTimes(2)
+        })
     })
-})
+}
 
 describe('userfriendlyImport', () => {
     it('returns nothing if pkg is empty', async () => {

tests/typings/webdriverio/globalImport.ts

@@ -1,6 +1,6 @@
 import { expectType } from 'tsd'
 import { $, $$, browser, driver, multiremotebrowser } from '@wdio/globals'
-import { fn, spyOn } from '@wdio/browser-runner'
+import { fn, spyOn, mock, unmock, mocked } from '@wdio/browser-runner'
 
 ;(async () => {
     const elem = await $('foo')
@@ -23,4 +23,11 @@ import { fn, spyOn } from '@wdio/browser-runner'
 
     expectType<Function>(fn)
     expectType<Function>(spyOn)
+
+    mock('foobar', (importOrig) => {
+        expectType<Promise<unknown>>(importOrig())
+        return { default: 'foobar' }
+    })
+    unmock('foobar')
+    expectType<[string][]>(mocked((foo: string) => 'bar').mock.calls)
 })

vitest.config.mts

@@ -22,10 +22,10 @@ export default defineConfig({
                 'packages/devtools/src/commands',
                 'packages/devtools/src/scripts'
             ],
-            lines: 93,
-            functions: 89,
-            branches: 93,
-            statements: 93
+            lines: 92,
+            functions: 88,
+            branches: 92,
+            statements: 92
         },
         globalSetup: [
             'scripts/test/globalSetup.ts'