fix: work around getMockImplementation() in Vitest (#28)

Ensure that default implementations set via `vi.fn(default_impl)` are used,
even when Vitest forgets they exists, by pulling from tinyspy internals

Author: tommie-lie

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 = <TFunc extends AnyCallable>(
+  mock: MockInstance<TFunc>,
+): TFunc | undefined => {
+  return (
+    mock.getMockImplementation() ?? getTinyspyInternals(mock)?.getOriginal()
+  )
+}
+
+/** Internal state from Tinyspy, where a mock's default implementation is stored. */
+interface TinyspyInternals<TFunc extends AnyCallable> {
+  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 = <TFunc extends AnyCallable>(
+  mock: MockInstance<TFunc>,
+): TinyspyInternals<TFunc> | undefined => {
+  const maybeTinyspy = mock as unknown as Record<PropertyKey, unknown>
+
+  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<TFunc>
+    }
+  }
+
+  return undefined
+}

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 = <TFunc extends AnyCallable>(
   }
 
   const behaviors = createBehaviorStack<TFunc>()
-  const fallbackImplementation = spy.getMockImplementation()
+  const fallbackImplementation = getFallbackImplementation(spy)
 
   const implementation = (...args: ExtractParameters<TFunc>) => {
     const behavior = behaviors.use(args)?.behavior ?? {

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()