diff --git a/src/fallback-implementation.ts b/src/fallback-implementation.ts new file mode 100644 index 0000000..9e915be --- /dev/null +++ b/src/fallback-implementation.ts @@ -0,0 +1,45 @@ +import type { AnyCallable, MockInstance } from './types.ts' + +/** Get the fallback implementation of a mock if no matching stub is found. */ +export const getFallbackImplementation = ( + mock: MockInstance, +): TFunc | undefined => { + return ( + mock.getMockImplementation() ?? getTinyspyInternals(mock)?.getOriginal() + ) +} + +/** Internal state from Tinyspy, where a mock's default implementation is stored. */ +interface TinyspyInternals { + getOriginal: () => TFunc | undefined +} + +/** + * Get the fallback implementation out of tinyspy internals. + * + * This slight hack works around a bug in Vitest <= 3 + * where `getMockImplementation` will return `undefined` after `mockReset`, + * even if a default implementation is still active. + * The implementation remains present in tinyspy internal state, + * which is stored on a Symbol key in the mock object. + */ +const getTinyspyInternals = ( + mock: MockInstance, +): TinyspyInternals | undefined => { + const maybeTinyspy = mock as unknown as Record + + for (const key of Object.getOwnPropertySymbols(maybeTinyspy)) { + const maybeTinyspyInternals = maybeTinyspy[key] + + if ( + maybeTinyspyInternals && + typeof maybeTinyspyInternals === 'object' && + 'getOriginal' in maybeTinyspyInternals && + typeof maybeTinyspyInternals.getOriginal === 'function' + ) { + return maybeTinyspyInternals as TinyspyInternals + } + } + + return undefined +} diff --git a/src/stubs.ts b/src/stubs.ts index 77417ae..fe92e1d 100644 --- a/src/stubs.ts +++ b/src/stubs.ts @@ -4,6 +4,7 @@ import { createBehaviorStack, } from './behaviors.ts' import { NotAMockFunctionError } from './errors.ts' +import { getFallbackImplementation } from './fallback-implementation.ts' import type { AnyCallable, AnyFunction, @@ -29,7 +30,7 @@ export const configureStub = ( } const behaviors = createBehaviorStack() - const fallbackImplementation = spy.getMockImplementation() + const fallbackImplementation = getFallbackImplementation(spy) const implementation = (...args: ExtractParameters) => { const behavior = behaviors.use(args)?.behavior ?? { diff --git a/test/vitest-when.test.ts b/test/vitest-when.test.ts index 676c236..35c6524 100644 --- a/test/vitest-when.test.ts +++ b/test/vitest-when.test.ts @@ -64,6 +64,17 @@ describe('vitest-when', () => { expect(spy()).toEqual(100) }) + it('should fall back to original implementation after reset', () => { + const spy = vi.fn((n) => 2 * n) + + vi.resetAllMocks() + expect(spy(2)).toEqual(4) + + subject.when(spy).calledWith(1).thenReturn(4) + expect(spy(1)).toEqual(4) + expect(spy(2)).toEqual(4) + }) + it('should return a number of times', () => { const spy = vi.fn()