fix(types): accept MockInstance as stubbing target

src/stubs.ts

@@ -1,11 +1,10 @@
-import type { Mock } from 'vitest'
 import {
   createBehaviorStack,
   type BehaviorStack,
   BehaviorType,
 } from './behaviors.ts'
 import { NotAMockFunctionError } from './errors.ts'
-import type { AnyFunction } from './types.ts'
+import type { AnyFunction, MockInstance } from './types.ts'
 
 const BEHAVIORS_KEY = Symbol('behaviors')
 
@@ -63,7 +62,7 @@ export const configureStub = <TFunc extends AnyFunction>(
   return behaviors
 }
 
-export const validateSpy = (maybeSpy: unknown): Mock => {
+export const validateSpy = (maybeSpy: unknown): MockInstance => {
   if (
     typeof maybeSpy === 'function' &&
     'mockImplementation' in maybeSpy &&
@@ -73,14 +72,14 @@ export const validateSpy = (maybeSpy: unknown): Mock => {
     'getMockName' in maybeSpy &&
     typeof maybeSpy.getMockName === 'function'
   ) {
-    return maybeSpy as Mock
+    return maybeSpy as unknown as MockInstance
   }
 
   throw new NotAMockFunctionError(maybeSpy)
 }
 
 export const getBehaviorStack = <TFunc extends AnyFunction>(
-  spy: Mock,
+  spy: MockInstance,
 ): BehaviorStack<TFunc> | undefined => {
   const existingImplementation = spy.getMockImplementation() as
     | WhenStubImplementation<TFunc>

src/types.ts

@@ -2,3 +2,20 @@
 
 /** Any function, for use in `extends` */
 export type AnyFunction = (...args: never[]) => unknown
+
+/**
+ * Minimally typed version of Vitest's `MockInstance`.
+ *
+ * Used to ensure backwards compatibility
+ * with older versions of Vitest.
+ */
+export interface MockInstance<TFunc extends AnyFunction = AnyFunction> {
+  getMockName(): string
+  getMockImplementation(): TFunc | undefined
+  mockImplementation: (impl: TFunc) => this
+  mock: MockContext<TFunc>
+}
+
+export interface MockContext<TFunc extends AnyFunction> {
+  calls: Parameters<TFunc>[]
+}

src/vitest-when.ts

@@ -1,6 +1,6 @@
 import { configureStub } from './stubs.ts'
 import type { WhenOptions } from './behaviors.ts'
-import type { AnyFunction } from './types.ts'
+import type { AnyFunction, MockInstance } from './types.ts'
 import { getDebug, type DebugResult } from './debug.ts'
 
 export { type WhenOptions, type Behavior, BehaviorType } from './behaviors.ts'
@@ -22,7 +22,7 @@ export interface Stub<TArgs extends unknown[], TReturn> {
 }
 
 export const when = <TFunc extends AnyFunction>(
-  spy: TFunc,
+  spy: TFunc | MockInstance<TFunc>,
   options: WhenOptions = {},
 ): StubWrapper<TFunc> => {
   const behaviorStack = configureStub(spy)

test/typing.test-d.ts

@@ -23,6 +23,16 @@ describe('vitest-when type signatures', () => {
     assertType<subject.Stub<[number], any>>(stub)
   })
 
+  it('should handle an spied function', () => {
+    const target = { simple }
+    const spy = vi.spyOn(target, 'simple')
+    const stub = subject.when(spy).calledWith(1)
+
+    stub.thenReturn('hello')
+
+    assertType<subject.Stub<[number], any>>(stub)
+  })
+
   it('should handle a simple function', () => {
     const stub = subject.when(simple).calledWith(1)