fix: resolve id before loading from vite (#1204)

* fix: resolve id before loading from vite * fix: allow import actual modules inside a factory of __mocks__ file * test: fix custom reporters test * fix: invalidate mocks * chore: fix compatability with node 14 * chore: fix callstack position * chore: remove collect from worker.ts since entry doesn't export that function * chore: fix pnpm-lock * chore: slice "mock" from stacktrace * chore: fix checking callstack for a function mock * chore: add description to factory mock * test: add test for importing itself inside __mocks__ folder
vitest-dev · May 7, 2022 · 110788a · 110788a
1 parent fbd7974
commit 110788a
Show file tree

Hide file tree

Showing 14 changed files with 136 additions and 37 deletions.
diff --git a/examples/mocks/package.json b/examples/mocks/package.json
@@ -19,7 +19,8 @@
   "devDependencies": {
     "@vitest/ui": "latest",
     "vite": "^2.9.5",
-    "vitest": "latest"
+    "vitest": "latest",
+    "zustand": "4.0.0-rc.1"
   },
   "stackblitz": {
     "startCommand": "npm run test:ui"

diff --git a/examples/mocks/src/zustand.ts b/examples/mocks/src/zustand.ts
@@ -0,0 +1,14 @@
+import actualCreate from 'zustand'
+
+// a variable to hold reset functions for all stores declared in the app
+const storeResetFns = new Set()
+
+// when creating a store, we get its initial state, create a reset function and add it in the set
+const create = vi.fn((createState) => {
+  const store = actualCreate(createState)
+  const initialState = store.getState()
+  storeResetFns.add(() => store.setState(initialState, true))
+  return store
+})
+
+export default create
diff --git a/examples/mocks/test/factory.test.ts b/examples/mocks/test/factory.test.ts
@@ -1,6 +1,7 @@
 import axios from 'axios'
 import * as example from '../src/example'
 import * as moduleA from '../src/moduleA'
+import logger from '../src/log'
 
 vi
   .mock('../src/example', () => ({
@@ -28,6 +29,17 @@ vi.mock('axios', () => {
   }
 })
 
+vi.mock('../src/log.ts', async () => {
+  // can import the same module inside and does not go into an infinite loop
+  const log = await import('../src/log')
+  return {
+    default: {
+      ...log.default,
+      info: vi.fn(),
+    },
+  }
+})
+
 describe('mocking with factory', () => {
   test('successfuly mocked', () => {
     expect((example as any).mocked).toBe(true)
@@ -44,4 +56,10 @@ describe('mocking with factory', () => {
 
     expect(axios.get).toHaveBeenCalledTimes(1)
   })
+
+  test('logger extended', () => {
+    expect(logger.warn).toBeTypeOf('function')
+    // @ts-expect-error extending module
+    expect(logger.info).toBeTypeOf('function')
+  })
 })
diff --git a/examples/mocks/test/self-importing.test.ts b/examples/mocks/test/self-importing.test.ts
@@ -0,0 +1,9 @@
+import zustand from 'zustand'
+
+vi.mock('zustand')
+
+describe('zustand didn\'t go into an infinite loop', () => {
+  test('zustand is mocked', () => {
+    expect(vi.isMockFunction(zustand)).toBe(true)
+  })
+})
diff --git a/packages/vite-node/src/client.ts b/packages/vite-node/src/client.ts
@@ -95,13 +95,6 @@ export class ViteNodeRunner {
         return `stack:\n${[...callstack, dep].reverse().map(p => `- ${p}`).join('\n')}`
       }
 
-      // probably means it was passed as variable
-      // and wasn't transformed by Vite
-      if (this.options.resolveId && this.shouldResolveId(dep)) {
-        const resolvedDep = await this.options.resolveId(dep, id)
-        dep = resolvedDep?.id?.replace(this.root, '') || dep
-      }
-
       let debugTimer: any
       if (this.debug)
         debugTimer = setTimeout(() => this.debugLog(() => `module ${dep} takes over 2s to load.\n${getStack()}`), 2000)
@@ -125,6 +118,27 @@ export class ViteNodeRunner {
       }
     }
 
+    Object.defineProperty(request, 'callstack', { get: () => callstack })
+
+    const resolveId = async (dep: string, callstackPosition = 1) => {
+      // probably means it was passed as variable
+      // and wasn't transformed by Vite
+      // or some dependency name was passed
+      // runner.executeFile('@scope/name')
+      // runner.executeFile(myDynamicName)
+      if (this.options.resolveId && this.shouldResolveId(dep)) {
+        let importer = callstack[callstack.length - callstackPosition]
+        if (importer && importer.startsWith('mock:'))
+          importer = importer.slice(5)
+        const { id } = await this.options.resolveId(dep, importer) || {}
+        dep = id && isAbsolute(id) ? `/@fs/${id}` : id || dep
+      }
+
+      return dep
+    }
+
+    id = await resolveId(id, 2)
+
     const requestStubs = this.options.requestStubs || DEFAULT_REQUEST_STUBS
     if (id in requestStubs)
       return requestStubs[id]
@@ -158,6 +172,8 @@ export class ViteNodeRunner {
       },
     }
 
+    let require: NodeRequire
+
     // Be careful when changing this
     // changing context will change amount of code added on line :114 (vm.runInThisContext)
     // this messes up sourcemaps for coverage
@@ -170,8 +186,15 @@ export class ViteNodeRunner {
       __vite_ssr_exportAll__: (obj: any) => exportAll(exports, obj),
       __vite_ssr_import_meta__: { url },
 
+      __vitest_resolve_id__: resolveId,
+
       // cjs compact
-      require: createRequire(url),
+      require: (path: string) => {
+        if (!require)
+          require = createRequire(url)
+
+        return require(path)
+      },
       exports,
       module: moduleProxy,
       __filename,
@@ -197,7 +220,7 @@ export class ViteNodeRunner {
   }
 
   shouldResolveId(dep: string) {
-    if (isNodeBuiltin(dep) || dep in (this.options.requestStubs || DEFAULT_REQUEST_STUBS))
+    if (isNodeBuiltin(dep) || dep in (this.options.requestStubs || DEFAULT_REQUEST_STUBS) || dep.startsWith('/@vite'))
       return false
 
     return !isAbsolute(dep) || !extname(dep)

diff --git a/packages/vitest/src/integrations/coverage.ts b/packages/vitest/src/integrations/coverage.ts
@@ -83,7 +83,7 @@ export async function reportCoverage(ctx: Vitest) {
   // This is a magic number. It corresponds to the amount of code
   // that we add in packages/vite-node/src/client.ts:114 (vm.runInThisContext)
   // TODO: Include our transformations in sourcemaps
-  const offset = 203
+  const offset = 224
 
   report._getSourceMap = (coverage: Profiler.ScriptCoverage) => {
     const path = pathToFileURL(coverage.url).href

diff --git a/packages/vitest/src/integrations/spy.ts b/packages/vitest/src/integrations/spy.ts
@@ -167,7 +167,7 @@ function enhanceSpy<TArgs extends any[], TReturns>(
       })
     },
     get lastCall() {
-      return stub.calls.at(-1)
+      return stub.calls[stub.calls.length - 1]
     },
   }
 

diff --git a/packages/vitest/src/node/core.ts b/packages/vitest/src/node/core.ts
@@ -78,12 +78,12 @@ export class Vitest {
       fetchModule(id: string) {
         return node.fetchModule(id)
       },
-      resolveId(id: string, importer: string | undefined) {
+      resolveId(id: string, importer?: string) {
         return node.resolveId(id, importer)
       },
     })
 
-    this.reporters = await createReporters(resolved.reporters, this.runner.executeFile.bind(this.runner))
+    this.reporters = await createReporters(resolved.reporters, this.runner)
 
     this.runningPromise = undefined
 

diff --git a/packages/vitest/src/node/reporters/utils.ts b/packages/vitest/src/node/reporters/utils.ts
@@ -1,11 +1,12 @@
+import type { ViteNodeRunner } from 'vite-node/client'
 import type { Reporter } from '../../types'
 import { ReportersMap } from './index'
 import type { BuiltinReporters } from './index'
 
-async function loadCustomReporterModule<C extends Reporter>(path: string, fetchModule: (id: string) => Promise<any>): Promise<new () => C> {
+async function loadCustomReporterModule<C extends Reporter>(path: string, runner: ViteNodeRunner): Promise<new () => C> {
   let customReporterModule: { default: new () => C }
   try {
-    customReporterModule = await fetchModule(path)
+    customReporterModule = await runner.executeId(path)
   }
   catch (customReporterModuleError) {
     throw new Error(`Failed to load custom Reporter from ${path}`, { cause: customReporterModuleError as Error })
@@ -17,15 +18,15 @@ async function loadCustomReporterModule<C extends Reporter>(path: string, fetchM
   return customReporterModule.default
 }
 
-function createReporters(reporterReferences: Array<string | Reporter | BuiltinReporters>, fetchModule: (id: string) => Promise<any>) {
+function createReporters(reporterReferences: Array<string | Reporter | BuiltinReporters>, runner: ViteNodeRunner) {
   const promisedReporters = reporterReferences.map(async (referenceOrInstance) => {
     if (typeof referenceOrInstance === 'string') {
       if (referenceOrInstance in ReportersMap) {
         const BuiltinReporter = ReportersMap[referenceOrInstance as BuiltinReporters]
         return new BuiltinReporter()
       }
       else {
-        const CustomReporter = await loadCustomReporterModule(referenceOrInstance, fetchModule)
+        const CustomReporter = await loadCustomReporterModule(referenceOrInstance, runner)
         return new CustomReporter()
       }
     }

diff --git a/packages/vitest/src/runtime/execute.ts b/packages/vitest/src/runtime/execute.ts
@@ -33,6 +33,7 @@ export class VitestRunner extends ViteNodeRunner {
 
   prepareContext(context: Record<string, any>) {
     const request = context.__vite_ssr_import__
+    const resolveId = context.__vitest_resolve_id__
 
     const mocker = this.mocker.withRequest(request)
 
@@ -49,8 +50,8 @@ export class VitestRunner extends ViteNodeRunner {
     }
 
     return Object.assign(context, {
-      __vite_ssr_import__: (dep: string) => mocker.requestWithMock(dep),
-      __vite_ssr_dynamic_import__: (dep: string) => mocker.requestWithMock(dep),
+      __vite_ssr_import__: async (dep: string) => mocker.requestWithMock(await resolveId(dep)),
+      __vite_ssr_dynamic_import__: async (dep: string) => mocker.requestWithMock(await resolveId(dep)),
       __vitest_mocker__: mocker,
     })
   }

diff --git a/packages/vitest/src/runtime/mocker.ts b/packages/vitest/src/runtime/mocker.ts
@@ -10,19 +10,24 @@ import type { ExecuteOptions } from './execute'
 
 type Callback = (...args: any[]) => unknown
 
+interface ViteRunnerRequest {
+  (dep: string): any
+  callstack: string[]
+}
+
 export class VitestMocker {
   private static pendingIds: PendingSuiteMock[] = []
   private static spyModule?: typeof import('../integrations/spy')
 
-  private request!: (dep: string) => unknown
+  private request!: ViteRunnerRequest
 
   private root: string
   private callbacks: Record<string, ((...args: any[]) => unknown)[]> = {}
 
   constructor(
     public options: ExecuteOptions,
     private moduleCache: ModuleCacheMap,
-    request?: (dep: string) => unknown,
+    request?: ViteRunnerRequest,
   ) {
     this.root = this.options.root
     this.request = request!
@@ -230,6 +235,8 @@ export class VitestMocker {
 
     const mock = this.getDependencyMock(dep)
 
+    const callstack = this.request.callstack
+
     if (mock === null) {
       const cacheName = `${dep}__mock`
       const cache = this.moduleCache.get(cacheName)
@@ -241,9 +248,11 @@ export class VitestMocker {
       this.emit('mocked', cacheName, { exports })
       return exports
     }
-    if (typeof mock === 'function')
+    if (typeof mock === 'function' && !callstack.includes(`mock:${dep}`)) {
+      callstack.push(`mock:${dep}`)
       return this.callFunctionMock(dep, mock)
-    if (typeof mock === 'string')
+    }
+    if (typeof mock === 'string' && !callstack.includes(mock))
       dep = mock
     return this.request(dep)
   }
@@ -256,7 +265,7 @@ export class VitestMocker {
     VitestMocker.pendingIds.push({ type: 'unmock', id, importer })
   }
 
-  public withRequest(request: (dep: string) => unknown) {
+  public withRequest(request: ViteRunnerRequest) {
     return new VitestMocker(this.options, this.moduleCache, request)
   }
 }
diff --git a/packages/vitest/src/runtime/worker.ts b/packages/vitest/src/runtime/worker.ts
@@ -10,7 +10,6 @@ import { rpc } from './rpc'
 
 let _viteNode: {
   run: (files: string[], config: ResolvedConfig) => Promise<void>
-  collect: (files: string[], config: ResolvedConfig) => Promise<void>
 }
 
 const moduleCache = new ModuleCacheMap()
@@ -37,7 +36,7 @@ async function startViteNode(ctx: WorkerContext) {
 
   const { config } = ctx
 
-  const { run, collect } = (await executeInViteNode({
+  const { run } = (await executeInViteNode({
     files: [
       resolve(distDir, 'entry.js'),
     ],
@@ -54,7 +53,7 @@ async function startViteNode(ctx: WorkerContext) {
     base: config.base,
   }))[0]
 
-  _viteNode = { run, collect }
+  _viteNode = { run }
 
   return _viteNode
 }
@@ -86,17 +85,15 @@ function init(ctx: WorkerContext) {
     ),
   }
 
-  if (ctx.invalidates)
-    ctx.invalidates.forEach(i => moduleCache.delete(i))
+  if (ctx.invalidates) {
+    ctx.invalidates.forEach((i) => {
+      moduleCache.delete(i)
+      moduleCache.delete(`${i}__mock`)
+    })
+  }
   ctx.files.forEach(i => moduleCache.delete(i))
 }
 
-export async function collect(ctx: WorkerContext) {
-  init(ctx)
-  const { collect } = await startViteNode(ctx)
-  return collect(ctx.files, ctx.config)
-}
-
 export async function run(ctx: WorkerContext) {
   init(ctx)
   const { run } = await startViteNode(ctx)

diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml