From a465428304115b27e072eebc12cbd1c8c03b3a94 Mon Sep 17 00:00:00 2001 From: Simon Abbott Date: Fri, 15 Jul 2022 15:07:49 -0500 Subject: [PATCH] fix: mock properties of classes/functions (fix #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 #1523. --- packages/vitest/src/runtime/mocker.ts | 119 ++++++++++++++++++++------ test/core/src/mockedC.ts | 24 ++++++ test/core/test/mocked.test.ts | 40 ++++++++- 3 files changed, 152 insertions(+), 31 deletions(-) create mode 100644 test/core/src/mockedC.ts diff --git a/packages/vitest/src/runtime/mocker.ts b/packages/vitest/src/runtime/mocker.ts index f7a65f7589c2..89a00f25513c 100644 --- a/packages/vitest/src/runtime/mocker.ts +++ b/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() + private mockedValueMap = new Map() + + 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) { 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 = {} - - 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, newContainer: Record) => { + 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 = {} + 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 } diff --git a/test/core/src/mockedC.ts b/test/core/src/mockedC.ts new file mode 100644 index 000000000000..36f47ebcb0ee --- /dev/null +++ b/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 { + await new Promise(resolve => resolve()) + return '1234' +} diff --git a/test/core/test/mocked.test.ts b/test/core/test/mocked.test.ts index 6880b5179a03..8025f52b8818 100644 --- a/test/core/test/mocked.test.ts +++ b/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') +})