diff --git a/packages/spy/src/index.ts b/packages/spy/src/index.ts index c271e7bc8..c3000c6cb 100644 --- a/packages/spy/src/index.ts +++ b/packages/spy/src/index.ts @@ -19,18 +19,55 @@ export function isMockFunction(fn: any): fn is Mock { } const MOCK_RESTORE = new Set<() => void>() +type MockRef = WeakRef> // Jest keeps the state in a separate WeakMap which is good for memory, // but it makes the state slower to access and return different values // if you stored it before calling `mockClear` where it will be recreated -const DIRTY_MOCK_STATES = new Set>() -const addDirtyMock = DIRTY_MOCK_STATES.add.bind(DIRTY_MOCK_STATES) -const REGISTERED_MOCKS = new Set>>() -const MOCK_FINALIZER = new FinalizationRegistry>>((ref) => { - REGISTERED_MOCKS.delete(ref) +const DIRTY_MOCK_STATES = new Set() +// Mocks whose implementation, once-queue or name changed since their last +// `mockReset()`, so `resetAllMocks()` only visits mocks that are not already reset +const DIRTY_MOCK_CONFIGS = new Set() +// a test may spy on `Set.prototype.add`, and the spy itself marks mocks dirty +const addToSet = Set.prototype.add +const MOCK_REFS = new WeakMap, MockRef>() +const MOCK_FINALIZER = new FinalizationRegistry((ref) => { + DIRTY_MOCK_STATES.delete(ref) + DIRTY_MOCK_CONFIGS.delete(ref) }) const MOCK_CONFIGS = new WeakMap, MockConfig>() const MOCKS_BY_STATE = new WeakMap>() +function markDirty(set: Set, mock: Mock): void { + let ref = MOCK_REFS.get(mock) + if (!ref) { + ref = new WeakRef(mock) + MOCK_REFS.set(mock, ref) + MOCK_FINALIZER.register(mock, ref) + } + addToSet.call(set, ref) +} + +function unmarkDirty(set: Set, mock: Mock): void { + const ref = MOCK_REFS.get(mock) + if (ref) { + set.delete(ref) + } +} + +function forEachDirty( + set: Set, + fn: (mock: Mock) => void, +): void { + for (const ref of set) { + const mock = ref.deref() + if (mock) { + fn(mock) + } else { + set.delete(ref) + } + } +} + export function createMockInstance( options: MockInstanceOption = {}, ): Mock { @@ -69,9 +106,6 @@ export function createMockInstance( } MOCK_CONFIGS.set(mock, config) MOCKS_BY_STATE.set(state, mock) - const ref = new WeakRef(mock) - REGISTERED_MOCKS.add(ref) - MOCK_FINALIZER.register(mock, ref) mock._isMockFunction = true mock.getMockImplementation = () => { @@ -101,24 +135,28 @@ export function createMockInstance( }) mock.mockImplementation = function mockImplementation(implementation) { + markDirty(DIRTY_MOCK_CONFIGS, mock) config.mockImplementation = implementation updateMockPrototype() return mock } mock.mockImplementationOnce = function mockImplementationOnce(implementation) { + markDirty(DIRTY_MOCK_CONFIGS, mock) config.onceMockImplementations.push(implementation) updateMockPrototype() return mock } mock.withImplementation = function withImplementation(implementation, callback) { + markDirty(DIRTY_MOCK_CONFIGS, mock) const previousImplementation = config.mockImplementation const previousOnceImplementations = config.onceMockImplementations const reset = () => { config.mockImplementation = previousImplementation config.onceMockImplementations = previousOnceImplementations + markDirty(DIRTY_MOCK_CONFIGS, mock) updateMockPrototype() } @@ -229,7 +267,7 @@ export function createMockInstance( state.invocationCallOrder = [] state.results = [] state.settledResults = [] - DIRTY_MOCK_STATES.delete(mock) + unmarkDirty(DIRTY_MOCK_STATES, mock) return mock } @@ -238,6 +276,7 @@ export function createMockInstance( config.mockImplementation = resetToMockImplementation ? mockImplementation : undefined config.mockName = resetToMockName ? mock.name || 'vi.fn()' : 'vi.fn()' config.onceMockImplementations = [] + unmarkDirty(DIRTY_MOCK_CONFIGS, mock) updateMockPrototype() return mock } @@ -249,6 +288,7 @@ export function createMockInstance( mock.mockName = function mockName(name: string) { if (typeof name === 'string') { + markDirty(DIRTY_MOCK_CONFIGS, mock) config.mockName = name } return mock @@ -484,9 +524,9 @@ function createMock({ const namedObject: Record> = { // to keep the name of the function intact [name]: function (this: any, ...args: any[]) { - addDirtyMock(namedObject[name]) + markDirty(DIRTY_MOCK_STATES, namedObject[name]) if (prototypeMock) { - addDirtyMock(prototypeMock) + markDirty(DIRTY_MOCK_STATES, prototypeMock) } registerCalls(args, state, prototypeState) registerInvocationOrder(invocationCallCounter++, state, prototypeState) @@ -773,20 +813,14 @@ export function restoreAllMocks(): void { } export function clearAllMocks(): void { - for (const mock of DIRTY_MOCK_STATES) { - mock.mockClear() - } + forEachDirty(DIRTY_MOCK_STATES, (mock) => mock.mockClear()) } export function resetAllMocks(): void { - for (const ref of REGISTERED_MOCKS) { - const mock = ref.deref() - if (mock) { - mock.mockReset() - } else { - REGISTERED_MOCKS.delete(ref) - } - } + // `mockReset()` removes the mock from both sets, so a mock present in both + // is visited once. + forEachDirty(DIRTY_MOCK_STATES, (mock) => mock.mockReset()) + forEachDirty(DIRTY_MOCK_CONFIGS, (mock) => mock.mockReset()) } function throwConstructorError(shorthand: string): never { diff --git a/test/unit/test/mocking/vi-fn-retention.test.ts b/test/unit/test/mocking/vi-fn-retention.test.ts new file mode 100644 index 000000000..8afa7cdd1 --- /dev/null +++ b/test/unit/test/mocking/vi-fn-retention.test.ts @@ -0,0 +1,63 @@ +import type { Mock } from 'vitest' +import v8 from 'node:v8' +import vm from 'node:vm' +import { expect, test, vi } from 'vitest' + +// node exposes no public gc hook, but the flag can be flipped just long +// enough to grab one from a throwaway context +v8.setFlagsFromString('--expose-gc') +const gc = vm.runInNewContext('gc') +v8.setFlagsFromString('--no-expose-gc') + +type Fn = Mock<() => unknown> + +function createMocks(create: () => Fn, use: (mock: Fn) => void): WeakRef[] { + return Array.from({ length: 100 }, () => { + const mock = create() + use(mock) + return new WeakRef(mock) + }) +} + +async function countAlive(refs: WeakRef[]): Promise { + let alive = refs.length + for (let attempt = 0; attempt < 10 && alive > 0; attempt++) { + await new Promise((resolve) => setTimeout(resolve, 10)) + gc() + alive = refs.filter((ref) => ref.deref()).length + } + return alive +} + +function noop() {} + +test.each([ + ['called, nothing clears them', () => vi.fn(), (mock: Fn) => mock(), noop], + ['configured, nothing clears them', () => vi.fn(), (mock: Fn) => mock.mockReturnValue(1), noop], + [ + 'created with an implementation, only vi.clearAllMocks() runs', + () => vi.fn(() => 1), + noop, + () => vi.clearAllMocks(), + ], + [ + 'configured, only vi.clearAllMocks() runs', + () => vi.fn(), + (mock: Fn) => mock.mockReturnValue(1), + () => vi.clearAllMocks(), + ], + [ + 'called and configured, only vi.clearAllMocks() runs', + () => vi.fn(), + (mock: Fn) => { + mock.mockReturnValue(1) + mock() + }, + () => vi.clearAllMocks(), + ], +])('unreferenced mocks are collected: %s', async (_, create, use, cleanup) => { + const refs = createMocks(create, use) + cleanup() + + expect(await countAlive(refs)).toBe(0) +}) diff --git a/test/unit/test/mocking/vi-fn.test.ts b/test/unit/test/mocking/vi-fn.test.ts index c121d9344..c39d50588 100644 --- a/test/unit/test/mocking/vi-fn.test.ts +++ b/test/unit/test/mocking/vi-fn.test.ts @@ -350,6 +350,70 @@ describe('vi.fn() configuration', () => { expect(mock()).toBe(undefined) }) + test('vi.resetAllMocks() only resets mocks that were called or reconfigured', () => { + const mocks = Array.from({ length: 100 }, () => vi.fn()) + mocks[10]() + mocks[20].mockReturnValue(42) + mocks[30].mockReturnValueOnce(42) + mocks[40].mockName('named') + mocks[50]() + mocks[50].mockReturnValue(42) + + const reset: number[] = [] + for (const [index, mock] of mocks.entries()) { + const mockReset = mock.mockReset + mock.mockReset = function () { + reset.push(index) + return mockReset.call(this) + } + } + + vi.resetAllMocks() + + expect(reset.sort((a, b) => a - b)).toEqual([10, 20, 30, 40, 50]) + expect(mocks[20]()).toBe(undefined) + expect(mocks[30]()).toBe(undefined) + expect(mocks[40].getMockName()).toBe('vi.fn()') + }) + + test('vi.resetAllMocks() still resets a reconfigured mock after vi.clearAllMocks()', () => { + const mock = vi.fn().mockReturnValue(42) + + vi.clearAllMocks() + vi.resetAllMocks() + + expect(mock()).toBe(undefined) + }) + + test('vi.resetAllMocks() resets an implementation restored by withImplementation after mockReset()', () => { + const mock = vi.fn().mockImplementation(() => 1) + mock.withImplementation( + () => 2, + () => { + mock.mockReset() + }, + ) + expect(mock.getMockImplementation()?.()).toBe(1) + + vi.resetAllMocks() + + expect(mock.getMockImplementation()).toBe(undefined) + }) + + test('vi.resetAllMocks() resets a once implementation restored by withImplementation after mockReset()', () => { + const mock = vi.fn().mockImplementationOnce(() => 1) + mock.withImplementation( + () => 2, + () => { + mock.mockReset() + }, + ) + + vi.resetAllMocks() + + expect(mock.getMockImplementation()).toBe(undefined) + }) + test('vi.fn() resets the original mock implementation', () => { const mock = vi.fn(() => 42) expect(mock()).toBe(42)