fix(types): accept MockInstance (again)

src/behaviors.ts

@@ -1,8 +1,8 @@
 import { equals } from '@vitest/expect'
 
 import type {
-  AnyCallable,
   AnyFunction,
+  Mock,
   ParametersOf,
   ReturnTypeOf,
   WithMatchers,
@@ -12,27 +12,25 @@ export interface WhenOptions {
   times?: number
 }
 
-export interface BehaviorStack<TFunc extends AnyCallable> {
-  use: (
-    args: ParametersOf<TFunc>,
-  ) => BehaviorEntry<ParametersOf<TFunc>> | undefined
+export interface BehaviorStack<TParameters extends unknown[], TReturn> {
+  use: (args: TParameters) => BehaviorEntry<TParameters> | undefined
 
-  getAll: () => readonly BehaviorEntry<ParametersOf<TFunc>>[]
+  getAll: () => readonly BehaviorEntry<TParameters>[]
 
-  getUnmatchedCalls: () => readonly ParametersOf<TFunc>[]
+  getUnmatchedCalls: () => readonly TParameters[]
 
   bindArgs: (
-    args: WithMatchers<ParametersOf<TFunc>>,
+    args: WithMatchers<TParameters>,
     options: WhenOptions,
-  ) => BoundBehaviorStack<ReturnTypeOf<TFunc>>
+  ) => BoundBehaviorStack<TParameters, TReturn>
 }
 
-export interface BoundBehaviorStack<TReturn> {
+export interface BoundBehaviorStack<TParameters extends unknown[], TReturn> {
   addReturn: (values: TReturn[]) => void
   addResolve: (values: Awaited<TReturn>[]) => void
   addThrow: (values: unknown[]) => void
   addReject: (values: unknown[]) => void
-  addDo: (values: AnyFunction[]) => void
+  addDo: (values: ((...args: TParameters) => TReturn)[]) => void
 }
 
 export interface BehaviorEntry<TArgs extends unknown[]> {
@@ -62,11 +60,16 @@ export interface BehaviorOptions<TValue> {
   maxCallCount: number | undefined
 }
 
+export type BehaviorStackOf<TMock extends Mock> = BehaviorStack<
+  ParametersOf<TMock>,
+  ReturnTypeOf<TMock>
+>
+
 export const createBehaviorStack = <
-  TFunc extends AnyCallable,
->(): BehaviorStack<TFunc> => {
-  const behaviors: BehaviorEntry<ParametersOf<TFunc>>[] = []
-  const unmatchedCalls: ParametersOf<TFunc>[] = []
+  TMock extends Mock,
+>(): BehaviorStackOf<TMock> => {
+  const behaviors: BehaviorEntry<ParametersOf<TMock>>[] = []
+  const unmatchedCalls: ParametersOf<TMock>[] = []
 
   return {
     getAll: () => behaviors,

src/debug.ts

@@ -5,7 +5,7 @@ import {
 
 import { type Behavior, BehaviorType } from './behaviors.ts'
 import { getBehaviorStack } from './stubs.ts'
-import type { AnyCallable, Mock } from './types.ts'
+import type { Mock } from './types.ts'
 
 export interface DebugResult {
   name: string
@@ -20,9 +20,7 @@ export interface Stubbing {
   calls: readonly unknown[][]
 }
 
-export const getDebug = <TFunc extends AnyCallable>(
-  mock: Mock<TFunc>,
-): DebugResult => {
+export const getDebug = (mock: Mock): DebugResult => {
   const name = mock.getMockName()
   const behaviors = getBehaviorStack(mock)
   const unmatchedCalls = behaviors?.getUnmatchedCalls() ?? mock.mock.calls

src/fallback-implementation.ts

@@ -1,18 +1,18 @@
-import type { AnyCallable, AsFunction, Mock } from './types.ts'
+import type { AsFunction, Mock } from './types.ts'
 
 /** Get the fallback implementation of a mock if no matching stub is found. */
-export const getFallbackImplementation = <TFunc extends AnyCallable>(
-  mock: Mock<TFunc>,
-): AsFunction<TFunc> | undefined => {
+export const getFallbackImplementation = <TMock extends Mock>(
+  mock: TMock,
+): AsFunction<TMock> | undefined => {
   return (
-    (mock.getMockImplementation() as AsFunction<TFunc> | undefined) ??
+    (mock.getMockImplementation() as AsFunction<TMock> | undefined) ??
     getTinyspyInternals(mock)?.getOriginal()
   )
 }
 
 /** Internal state from Tinyspy, where a mock's default implementation is stored. */
-interface TinyspyInternals<TFunc extends AnyCallable> {
-  getOriginal: () => AsFunction<TFunc> | undefined
+interface TinyspyInternals<TMock extends Mock> {
+  getOriginal: () => AsFunction<TMock> | undefined
 }
 
 /**
@@ -24,9 +24,9 @@ interface TinyspyInternals<TFunc extends AnyCallable> {
  * 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: Mock<TFunc>,
-): TinyspyInternals<TFunc> | undefined => {
+const getTinyspyInternals = <TMock extends Mock>(
+  mock: TMock,
+): TinyspyInternals<TMock> | undefined => {
   const maybeTinyspy = mock as unknown as Record<PropertyKey, unknown>
 
   for (const key of Object.getOwnPropertySymbols(maybeTinyspy)) {
@@ -38,7 +38,7 @@ const getTinyspyInternals = <TFunc extends AnyCallable>(
       'getOriginal' in maybeTinyspyInternals &&
       typeof maybeTinyspyInternals.getOriginal === 'function'
     ) {
-      return maybeTinyspyInternals as TinyspyInternals<TFunc>
+      return maybeTinyspyInternals as TinyspyInternals<TMock>
     }
   }
 

src/stubs.ts

@@ -1,32 +1,39 @@
 import {
   type BehaviorStack,
+  type BehaviorStackOf,
   BehaviorType,
   createBehaviorStack,
 } from './behaviors.ts'
 import { NotAMockFunctionError } from './errors.ts'
 import { getFallbackImplementation } from './fallback-implementation.ts'
-import type { AnyCallable, Mock, ParametersOf } from './types.ts'
+import type {
+  AsFunction,
+  Mock,
+  MockSource,
+  ParametersOf,
+  ReturnTypeOf,
+} from './types.ts'
 
 const BEHAVIORS_KEY = Symbol.for('vitest-when:behaviors')
 
-interface WhenStubImplementation<TFunc extends AnyCallable> {
-  (...args: ParametersOf<TFunc>): unknown
-  [BEHAVIORS_KEY]: BehaviorStack<TFunc>
+interface WhenStubImplementation<TMock extends Mock> {
+  (...args: ParametersOf<TMock>): unknown
+  [BEHAVIORS_KEY]: BehaviorStack<ParametersOf<TMock>, ReturnTypeOf<TMock>>
 }
 
-export const configureMock = <TFunc extends AnyCallable>(
-  mock: Mock<TFunc>,
-): BehaviorStack<TFunc> => {
+export const configureMock = <TMock extends Mock>(
+  mock: TMock,
+): BehaviorStackOf<TMock> => {
   const existingBehaviorStack = getBehaviorStack(mock)
 
   if (existingBehaviorStack) {
     return existingBehaviorStack
   }
 
-  const behaviorStack = createBehaviorStack<TFunc>()
+  const behaviorStack = createBehaviorStack<TMock>()
   const fallbackImplementation = getFallbackImplementation(mock)
 
-  const implementation = (...args: ParametersOf<TFunc>) => {
+  const implementation = (...args: ParametersOf<TMock>) => {
     const behavior = behaviorStack.use(args)?.behavior ?? {
       type: BehaviorType.DO,
       callback: fallbackImplementation,
@@ -51,6 +58,7 @@ export const configureMock = <TFunc extends AnyCallable>(
       }
 
       case BehaviorType.DO: {
+        // eslint-disable-next-line @typescript-eslint/no-unsafe-return
         return behavior.callback?.(...args)
       }
     }
@@ -63,26 +71,26 @@ export const configureMock = <TFunc extends AnyCallable>(
   return behaviorStack
 }
 
-export const validateMock = <TFunc extends AnyCallable>(
-  maybeMock: TFunc | Mock<TFunc>,
-): Mock<TFunc> => {
+export const validateMock = <TSource extends MockSource>(
+  maybeMock: TSource,
+): Mock<TSource> => {
   if (
     typeof maybeMock === 'function' &&
     'mockImplementation' in maybeMock &&
     typeof maybeMock.mockImplementation === 'function'
   ) {
-    return maybeMock
+    return maybeMock as Mock<TSource>
   }
 
   throw new NotAMockFunctionError(maybeMock)
 }
 
-export const getBehaviorStack = <TFunc extends AnyCallable>(
-  mock: Mock<TFunc>,
-): BehaviorStack<TFunc> | undefined => {
+export const getBehaviorStack = <TMock extends Mock>(
+  mock: TMock,
+): BehaviorStackOf<TMock> | undefined => {
   const existingImplementation = mock.getMockImplementation() as
-    | WhenStubImplementation<TFunc>
-    | TFunc
+    | WhenStubImplementation<TMock>
+    | AsFunction<TMock>
     | undefined
 
   return existingImplementation && BEHAVIORS_KEY in existingImplementation

src/types.ts

@@ -1,45 +1,71 @@
+/* eslint-disable @typescript-eslint/no-explicit-any */
 /** Common type definitions. */
 import type { AsymmetricMatcher } from '@vitest/expect'
 import type { MockedClass, MockedFunction } from 'vitest'
 
 /** Any function. */
-export type AnyFunction = (...args: never[]) => unknown
+export type AnyFunction = (...args: any[]) => any
 
 /** Any constructor. */
-export type AnyConstructor = new (...args: never[]) => unknown
+export type AnyConstructor = new (...args: any[]) => any
 
-/** Any callable, for use in `extends` */
-export type AnyCallable = AnyFunction | AnyConstructor
+/**
+ * Minimally typed version of Vitest's `MockInstance`.
+ *
+ * Ensures backwards compatibility with vitest@<=1
+ */
+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>[]
+}
+
+/** A function, constructor, or `vi.spyOn` return that's been mocked. */
+export type MockSource = AnyFunction | AnyConstructor | MockInstance
 
 /** Extract parameters from either a function or constructor. */
-export type ParametersOf<TFunc extends AnyCallable> = TFunc extends new (
+export type ParametersOf<TMock extends MockSource> = TMock extends new (
   ...args: infer P
 ) => unknown
   ? P
-  : TFunc extends (...args: infer P) => unknown
+  : TMock extends (...args: infer P) => unknown
     ? P
-    : never
+    : TMock extends MockInstance<(...args: infer P) => unknown>
+      ? P
+      : never
 
 /** Extract return type from either a function or constructor */
-export type ReturnTypeOf<TFunc extends AnyCallable> = TFunc extends new (
+export type ReturnTypeOf<TMock extends MockSource> = TMock extends new (
   ...args: never[]
 ) => infer R
   ? R
-  : TFunc extends (...args: never[]) => infer R
+  : TMock extends (...args: any[]) => infer R
     ? R
-    : never
+    : TMock extends MockInstance<(...args: never[]) => infer R>
+      ? R
+      : never
 
-export type AsFunction<TFunc extends AnyCallable> = (
-  ...args: ParametersOf<TFunc>
-) => ReturnTypeOf<TFunc>
+/** Convert a function or constructor type into a function type. */
+export type AsFunction<TMock extends MockSource> = (
+  ...args: ParametersOf<TMock>
+) => ReturnTypeOf<TMock>
 
 /** Accept a value or an AsymmetricMatcher in an arguments array */
 export type WithMatchers<T extends unknown[]> = {
   [K in keyof T]: AsymmetricMatcher<unknown> | T[K]
 }
 
-export type Mock<TFunc extends AnyCallable> = TFunc extends AnyFunction
-  ? MockedFunction<TFunc>
-  : TFunc extends AnyConstructor
-    ? MockedClass<TFunc>
-    : never
+/** A mocked function or constructor. */
+export type Mock<TMock extends MockSource = MockSource> =
+  TMock extends MockInstance<infer TFunc>
+    ? MockedFunction<TFunc>
+    : TMock extends AnyFunction
+      ? MockedFunction<TMock>
+      : TMock extends AnyConstructor
+        ? MockedClass<TMock>
+        : never

src/vitest-when.ts

@@ -2,9 +2,9 @@ import type { WhenOptions } from './behaviors.ts'
 import { type DebugResult, getDebug } from './debug.ts'
 import { configureMock, validateMock } from './stubs.ts'
 import type {
-  AnyCallable,
   AsFunction,
   Mock,
+  MockSource,
   ParametersOf,
   ReturnTypeOf,
   WithMatchers,
@@ -14,24 +14,24 @@ export { type Behavior, BehaviorType, type WhenOptions } from './behaviors.ts'
 export type { DebugResult, Stubbing } from './debug.ts'
 export * from './errors.ts'
 
-export interface StubWrapper<TFunc extends AnyCallable> {
-  calledWith<TArgs extends ParametersOf<TFunc>>(
+export interface StubWrapper<TMock extends Mock> {
+  calledWith<TArgs extends ParametersOf<TMock>>(
     ...args: WithMatchers<TArgs>
-  ): Stub<TFunc>
+  ): Stub<TMock>
 }
 
-export interface Stub<TFunc extends AnyCallable> {
-  thenReturn: (...values: ReturnTypeOf<TFunc>[]) => Mock<TFunc>
-  thenResolve: (...values: Awaited<ReturnTypeOf<TFunc>>[]) => Mock<TFunc>
-  thenThrow: (...errors: unknown[]) => Mock<TFunc>
-  thenReject: (...errors: unknown[]) => Mock<TFunc>
-  thenDo: (...callbacks: AsFunction<TFunc>[]) => Mock<TFunc>
+export interface Stub<TMock extends Mock> {
+  thenReturn: (...values: ReturnTypeOf<TMock>[]) => TMock
+  thenResolve: (...values: Awaited<ReturnTypeOf<TMock>>[]) => TMock
+  thenThrow: (...errors: unknown[]) => TMock
+  thenReject: (...errors: unknown[]) => TMock
+  thenDo: (...callbacks: AsFunction<TMock>[]) => TMock
 }
 
-export const when = <TFunc extends AnyCallable>(
-  mock: TFunc | Mock<TFunc>,
+export const when = <TMockSource extends MockSource>(
+  mock: TMockSource,
   options: WhenOptions = {},
-): StubWrapper<TFunc> => {
+): StubWrapper<Mock<TMockSource>> => {
   const validatedMock = validateMock(mock)
   const behaviorStack = configureMock(validatedMock)
 
@@ -69,8 +69,8 @@ export interface DebugOptions {
   log?: boolean
 }
 
-export const debug = <TFunc extends AnyCallable>(
-  mock: TFunc | Mock<TFunc>,
+export const debug = (
+  mock: MockSource,
   options: DebugOptions = {},
 ): DebugResult => {
   const log = options.log ?? true

test/typing.test-d.ts

@@ -69,10 +69,9 @@ describe('vitest-when type signatures', () => {
   })
 
   it('should handle an spied function', () => {
-    const target = { simple }
-    vi.spyOn(target, 'simple')
+    const target = vi.spyOn({ simple }, 'simple')
 
-    const result = subject.when(target.simple).calledWith(1).thenReturn('hello')
+    const result = subject.when(target).calledWith(1).thenReturn('hello')
 
     expectTypeOf(result.mock.calls).toEqualTypeOf<[number][]>()
     expectTypeOf(result).parameters.toEqualTypeOf<[number]>()