fix: automocking prototype fix (#1083)

examples/mocks/src/log.ts

@@ -0,0 +1,13 @@
+const writeSym = Symbol('write')
+const proto = {
+  [writeSym]() {
+    return 'hello'
+  },
+}
+const logger = {
+  warn() {
+    this[writeSym]()
+  },
+}
+Object.setPrototypeOf(logger, proto)
+export default logger

examples/mocks/test/example-mock.test.ts → examples/mocks/test/automocking.spec.ts

@@ -1,5 +1,9 @@
 import type * as exampleModule from '../src/example'
 
+import log from '../src/log'
+
+vi.mock('../src/log')
+
 test('all mocked are valid', async() => {
   const example = await vi.importMock<typeof exampleModule>('../src/example')
 
@@ -35,3 +39,14 @@ test('all mocked are valid', async() => {
   expect(example.boolean).toEqual(true)
   expect(example.symbol).toEqual(Symbol.for('a.b.c'))
 })
+
+test('automock doesn\'t throw when unmocked', () => {
+  // logger uses symbols on its prototype
+  // we are checking here that after unmocking
+  // these symbols are accessible
+  expect(() => {
+    log.warn()
+    vi.restoreAllMocks()
+    log.warn()
+  }).not.toThrow()
+})

packages/vitest/src/runtime/mocker.ts

@@ -10,26 +10,21 @@ import type { ExecuteOptions } from './execute'
 
 type Callback = (...args: any[]) => unknown
 
-function getObjectType(value: unknown): string {
+function getType(value: unknown): string {
   return Object.prototype.toString.apply(value).slice(8, -1)
 }
 
-function mockPrototype(spyOn: typeof import('../integrations/jest-mock')['spyOn'], proto: any) {
-  if (!proto) return null
-
-  const newProto: any = {}
-
-  const protoDescr = Object.getOwnPropertyDescriptors(proto)
-
-  // eslint-disable-next-line no-restricted-syntax
-  for (const d in protoDescr) {
-    Object.defineProperty(newProto, d, protoDescr[d])
-
-    if (typeof protoDescr[d].value === 'function')
-      spyOn(newProto, d).mockImplementation(() => {})
-  }
-
-  return newProto
+function getAllProperties(obj: any) {
+  const allProps = new Set<string>()
+  let curr = obj
+  do {
+    // we don't need propterties from 'Object'
+    if (curr === Object.prototype) break
+    const props = Object.getOwnPropertyNames(curr)
+    props.forEach(prop => allProps.add(prop))
+  // eslint-disable-next-line no-cond-assign
+  } while (curr = Object.getPrototypeOf(curr))
+  return Array.from(allProps)
 }
 
 export class VitestMocker {
@@ -156,32 +151,39 @@ export class VitestMocker {
     return existsSync(fullPath) ? fullPath.replace(this.root, '') : null
   }
 
-  public mockObject(obj: any) {
-    if (!VitestMocker.spyModule)
-      throw new Error('Internal Vitest error: Spy function is not defined.')
+  public mockValue(value: 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 = getObjectType(obj)
+    const type = getType(value)
 
-    if (Array.isArray(obj))
+    if (Array.isArray(value))
       return []
     else if (type !== 'Object' && type !== 'Module')
-      return obj
+      return value
 
-    const newObj = { ...obj }
+    const newObj: any = {}
 
-    const proto = mockPrototype(VitestMocker.spyModule.spyOn, Object.getPrototypeOf(obj))
-    Object.setPrototypeOf(newObj, proto)
+    const proproperties = getAllProperties(value)
 
-    // eslint-disable-next-line no-restricted-syntax
-    for (const k in obj) {
-      newObj[k] = this.mockObject(obj[k])
-      const type = getObjectType(obj[k])
+    for (const k of proproperties) {
+      newObj[k] = this.mockValue(value[k])
+      const type = getType(value[k])
 
-      if (type.includes('Function') && !obj[k]._isMockFunction) {
-        VitestMocker.spyModule.spyOn(newObj, k).mockImplementation(() => {})
+      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
       }
     }
+
+    // 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
   }
 
@@ -225,7 +227,7 @@ export class VitestMocker {
       await this.ensureSpy()
       const fsPath = this.getFsPath(path, external)
       const mod = await this.request(fsPath)
-      return this.mockObject(mod)
+      return this.mockValue(mod)
     }
     if (typeof mock === 'function')
       return this.callFunctionMock(path, mock)
@@ -250,7 +252,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.mockObject(mod)
+      const exports = this.mockValue(mod)
       this.emit('mocked', cacheName, { exports })
       return exports
     }