fix: mock properties of classes/functions (vitest-dev#1523)

Previously classes weren't mocked properly because their prototypes were
not being copied, so they lost all their methods. This commit rewrites
the automocker to properly handle:

- Functions with properties (namely classes)
- Circular objects (also to handle classes, thanks JavaScript)

With these changes classes can now be properly automocked!

This fixes vitest-dev#1523.

packages/vitest/src/runtime/mocker.ts

@@ -8,6 +8,34 @@ import { distDir } from '../constants'
 import type { PendingSuiteMock } from '../types/mocker'
 import type { ExecuteOptions } from './execute'
 
+class RefTracker {
+  private idMap = new Map<any, number>()
+  private mockedValueMap = new Map<number, any>()
+
+  public getId(value: any) {
+    return this.idMap.get(value)
+  }
+
+  public getMockedValue(id: number) {
+    return this.mockedValueMap.get(id)
+  }
+
+  public track(originalValue: any, mockedValue: any): number {
+    const newId = this.idMap.size
+    this.idMap.set(originalValue, newId)
+    this.mockedValueMap.set(newId, mockedValue)
+    return newId
+  }
+}
+
+type Key = string | symbol
+
+function isSpecialProp(prop: Key, parentType: string) {
+  return parentType.includes('Function')
+      && typeof prop === 'string'
+      && ['arguments', 'callee', 'caller', 'length', 'name'].includes(prop)
+}
+
 interface ViteRunnerRequest {
   (dep: string): any
   callstack: string[]
@@ -129,40 +157,77 @@ export class VitestMocker {
     return existsSync(fullPath) ? fullPath : null
   }
 
-  public mockValue(value: any) {
+  public mockObject(object: Record<string | symbol, any>) {
     if (!VitestMocker.spyModule) {
       throw new Error(
         'Error: Spy module is not defined. '
         + 'This is likely an internal bug in Vitest. '
         + 'Please report it to https://github.com/vitest-dev/vitest/issues')
     }
-
-    const type = getType(value)
-
-    if (Array.isArray(value))
-      return []
-    else if (type !== 'Object' && type !== 'Module')
-      return value
-
-    const newObj: Record<string | symbol, any> = {}
-
-    const properties = getAllProperties(value)
-
-    for (const k of properties) {
-      newObj[k] = this.mockValue(value[k])
-      const type = getType(value[k])
-
-      if (type.includes('Function') && !value[k]._isMockFunction) {
-        VitestMocker.spyModule.spyOn(newObj, k).mockImplementation(() => undefined)
-        Object.defineProperty(newObj[k], 'length', { value: 0 }) // tinyspy retains length, but jest doesnt
+    const spyModule = VitestMocker.spyModule
+
+    const finalizers = new Array<() => void>()
+    const refs = new RefTracker()
+
+    const mockPropertiesOf = (container: Record<Key, any>, newContainer: Record<Key, any>) => {
+      const containerType = getType(container)
+      for (const property of getAllProperties(container)) {
+        // ES Modules define their exports as getters. We want to process those.
+        if (containerType !== 'Module') {
+          // TODO: Mock getters/setters somehow?
+          const descriptor = Object.getOwnPropertyDescriptor(container, property)
+          if (descriptor?.get || descriptor?.set)
+            continue
+        }
+
+        // Skip special read-only props, we don't want to mess with those.
+        if (isSpecialProp(property, containerType))
+          continue
+
+        const value = container[property]
+
+        // Special handling of references we've seen before to prevent infinite
+        // recursion in circular objects.
+        const refId = refs.getId(value)
+        if (refId) {
+          finalizers.push(() => newContainer[property] = refs.getMockedValue(refId))
+          continue
+        }
+
+        const type = getType(value)
+
+        if (Array.isArray(value)) {
+          newContainer[property] = []
+          continue
+        }
+
+        const isFunction = type.includes('Function') && typeof value === 'function'
+        if ((!isFunction || value.__isMockFunction) && type !== 'Object' && type !== 'Module') {
+          newContainer[property] = value
+          continue
+        }
+
+        newContainer[property] = isFunction ? value : {}
+
+        if (isFunction) {
+          spyModule.spyOn(newContainer, property).mockImplementation(() => undefined)
+          // tinyspy retains length, but jest doesn't.
+          Object.defineProperty(newContainer[property], 'length', { value: 0 })
+        }
+
+        refs.track(value, newContainer[property])
+        mockPropertiesOf(value, newContainer[property])
       }
     }
 
-    // should be defined after object, because it may contain
-    // special logic on getting/settings properties
-    // and we don't want to invoke it
-    Object.setPrototypeOf(newObj, Object.getPrototypeOf(value))
-    return newObj
+    const mockedObject: Record<Key, any> = {}
+    mockPropertiesOf(object, mockedObject)
+
+    // Plug together refs
+    for (const finalizer of finalizers)
+      finalizer()
+
+    return mockedObject
   }
 
   public unmockPath(path: string) {
@@ -205,7 +270,7 @@ export class VitestMocker {
     if (mock === null) {
       await this.ensureSpy()
       const mod = await this.request(fsPath)
-      return this.mockValue(mod)
+      return this.mockObject(mod)
     }
 
     if (typeof mock === 'function')
@@ -236,7 +301,7 @@ export class VitestMocker {
         return cache.exports
       const cacheKey = toFilePath(dep, this.root)
       const mod = this.moduleCache.get(cacheKey)?.exports || await this.request(dep)
-      const exports = this.mockValue(mod)
+      const exports = this.mockObject(mod)
       this.moduleCache.set(cacheName, { exports })
       return exports
     }

test/core/src/mockedC.ts

@@ -0,0 +1,24 @@
+import { mockedA } from './mockedA'
+
+export class MockedC {
+  public value: number
+
+  constructor() {
+    this.value = 42
+  }
+
+  public doSomething() {
+    return mockedA()
+  }
+
+  get getSetProp(): number {
+    return 123
+  }
+
+  set getSetProp(_val: number) {}
+}
+
+export async function asyncFunc(): Promise<string> {
+  await new Promise<void>(resolve => resolve())
+  return '1234'
+}

test/core/test/mocked.test.ts

@@ -1,21 +1,22 @@
-import { assert, expect, test, vi, vitest } from 'vitest'
+import { assert, describe, expect, test, vi, vitest } from 'vitest'
 // @ts-expect-error not typed module
 import { value as virtualValue } from 'virtual-module'
 import { two } from '../src/submodule'
 import * as mocked from '../src/mockedA'
 import { mockedB } from '../src/mockedB'
+import { MockedC, asyncFunc } from '../src/mockedC'
 import * as globalMock from '../src/global-mock'
 
 vitest.mock('../src/submodule')
-vitest.mock('virtual-module', () => {
-  return { value: 'mock' }
-})
+vitest.mock('virtual-module', () => ({ value: 'mock' }))
+vitest.mock('../src/mockedC')
 
 test('submodule is mocked to return "two" as 3', () => {
   assert.equal(3, two)
 })
 
 test('globally mocked files are mocked', () => {
+  // Mocked in setup.ts
   expect(globalMock.mocked).toBe(true)
 })
 
@@ -31,3 +32,34 @@ test('can mock esm', () => {
 test('mocked exports should override original exports', () => {
   expect(virtualValue).toBe('mock')
 })
+
+describe('mocked classes', () => {
+  test('should not delete the prototype', () => {
+    expect(MockedC).toBeTypeOf('function')
+    expect(MockedC.prototype.doSomething).toBeTypeOf('function')
+  })
+
+  test('should mock the constructor', () => {
+    const instance = new MockedC()
+
+    expect(instance.value).not.toBe(42)
+    expect(MockedC).toHaveBeenCalledOnce()
+  })
+
+  test('should mock functions in the prototype', () => {
+    const instance = new MockedC()
+
+    expect(instance.doSomething).toBeTypeOf('function')
+    expect(instance.doSomething()).not.toBe('A')
+
+    expect(MockedC.prototype.doSomething).toHaveBeenCalledOnce()
+    expect(MockedC.prototype.doSomething).not.toHaveReturnedWith('A')
+  })
+})
+
+test('async functions should be mocked', () => {
+  expect(asyncFunc()).toBeUndefined()
+  expect(vi.mocked(asyncFunc).mockResolvedValue).toBeDefined()
+  vi.mocked(asyncFunc).mockResolvedValue('foo')
+  expect(asyncFunc()).resolves.toBe('foo')
+})