diff --git a/README.md b/README.md
index 54edb05..293e889 100644
--- a/README.md
+++ b/README.md
@@ -69,13 +69,13 @@ You should call `vi.resetAllMocks()` in your suite's `afterEach` hook to remove
 
 [vitest's mock functions]: https://vitest.dev/api/mock.html
 [stubs]: https://en.wikipedia.org/wiki/Test_stub
-[when]: #whenspy-tfunc-stubwrappertfunc
-[called-with]: #calledwithargs-targs-stubtargs-treturn
-[then-return]: #thenreturnvalue-treturn
-[then-resolve]: #thenresolvevalue-treturn
-[then-throw]: #thenthrowerror-unknown
-[then-reject]: #thenrejecterror-unknown
-[then-do]: #thendocallback-args-targs--treturn
+[when]: #whenmock-tfunc-options-whenoptions-stubwrappertfunc
+[called-with]: #calledwithargs-parameterstfunc-stubtfunc
+[then-return]: #thenreturnvalue-treturn---mocktfunc
+[then-resolve]: #thenresolvevalue-treturn---mocktfunc
+[then-throw]: #thenthrowerror-unknown---mocktfunc
+[then-reject]: #thenrejecterror-unknown---mocktfunc
+[then-do]: #thendocallback-args-targs--treturn---mocktfunc
 
 ### Why not vanilla Vitest mocks?
 
@@ -184,7 +184,7 @@ export const calculateQuestion = async (answer: number): Promise<string> => {
 
 ## API
 
-### `when(spy: TFunc, options?: WhenOptions): StubWrapper<TFunc>`
+### `when(mock: TFunc, options?: WhenOptions): StubWrapper<TFunc>`
 
 Configures a `vi.fn()` or `vi.spyOn()` mock function to act as a vitest-when stub. Adds an implementation to the function that initially no-ops, and returns an API to configure behaviors for given arguments using [`.calledWith(...)`][called-with]
 
@@ -192,11 +192,11 @@ Configures a `vi.fn()` or `vi.spyOn()` mock function to act as a vitest-when stu
 import { vi } from 'vitest'
 import { when } from 'vitest-when'
 
-const spy = vi.fn()
+const mock = vi.fn()
 
-when(spy)
+when(mock)
 
-expect(spy()).toBe(undefined)
+expect(mock()).toBe(undefined)
 ```
 
 #### Options
@@ -209,36 +209,32 @@ import type { WhenOptions } from 'vitest-when'
 | ------- | ------- | ------- | -------------------------------------------------- |
 | `times` | N/A     | integer | Only trigger configured behavior a number of times |
 
-### `.calledWith(...args: TArgs): Stub<TArgs, TReturn>`
+### `.calledWith(...args: Parameters<TFunc>): Stub<TFunc>`
 
 Create a stub that matches a given set of arguments which you can configure with different behaviors using methods like [`.thenReturn(...)`][then-return].
 
 ```ts
-const spy = vi.fn()
+const mock = when(vi.fn()).calledWith('hello').thenReturn('world')
 
-when(spy).calledWith('hello').thenReturn('world')
-
-expect(spy('hello')).toEqual('world')
+expect(mock('hello')).toEqual('world')
 ```
 
 When a call to a mock uses arguments that match those given to `calledWith`, a configured behavior will be triggered. All arguments must match, but you can use Vitest's [asymmetric matchers][] to loosen the stubbing:
 
 ```ts
-const spy = vi.fn()
-
-when(spy).calledWith(expect.any(String)).thenReturn('world')
+const mock = when(vi.fn()).calledWith(expect.any(String)).thenReturn('world')
 
-expect(spy('hello')).toEqual('world')
-expect(spy('anything')).toEqual('world')
+expect(mock('hello')).toEqual('world')
+expect(mock('anything')).toEqual('world')
 ```
 
 If `calledWith` is used multiple times, the last configured stubbing will be used.
 
 ```ts
-when(spy).calledWith('hello').thenReturn('world')
-expect(spy('hello')).toEqual('world')
-when(spy).calledWith('hello').thenReturn('goodbye')
-expect(spy('hello')).toEqual('goodbye')
+when(mock).calledWith('hello').thenReturn('world')
+expect(mock('hello')).toEqual('world')
+when(mock).calledWith('hello').thenReturn('goodbye')
+expect(mock('hello')).toEqual('goodbye')
 ```
 
 [asymmetric matchers]: https://vitest.dev/api/expect.html#expect-anything
@@ -269,26 +265,24 @@ when<() => null>(overloaded).calledWith().thenReturn(null)
 By default, if arguments do not match, a vitest-when stub will no-op and return `undefined`. You can customize this fallback by configuring your own unconditional behavior on the mock using Vitest's built-in [mock API][].
 
 ```ts
-const spy = vi.fn().mockReturnValue('you messed up!')
-
-when(spy).calledWith('hello').thenReturn('world')
+const mock = when(vi.fn(() => 'you messed up!')))
+  .calledWith('hello')
+  .thenReturn('world')
 
-spy('hello') // "world"
-spy('jello') // "you messed up!"
+mock('hello') // "world"
+mock('jello') // "you messed up!"
 ```
 
 [mock API]: https://vitest.dev/api/mock.html
 
-### `.thenReturn(value: TReturn)`
+### `.thenReturn(value: TReturn) -> Mock<TFunc>`
 
 When the stubbing is satisfied, return `value`
 
 ```ts
-const spy = vi.fn()
-
-when(spy).calledWith('hello').thenReturn('world')
+const mock = when(vi.fn()).calledWith('hello').thenReturn('world')
 
-expect(spy('hello')).toEqual('world')
+expect(mock('hello')).toEqual('world')
 ```
 
 To only return a value once, use the `times` option.
@@ -296,36 +290,30 @@ To only return a value once, use the `times` option.
 ```ts
 import { when } from 'vitest-when'
 
-const spy = vi.fn()
-
-when(spy, { times: 1 }).calledWith('hello').thenReturn('world')
+const mock = when(vi.fn(), { times: 1 }).calledWith('hello').thenReturn('world')
 
-expect(spy('hello')).toEqual('world')
-expect(spy('hello')).toEqual(undefined)
+expect(mock('hello')).toEqual('world')
+expect(mock('hello')).toEqual(undefined)
 ```
 
 You may pass several values to `thenReturn` to return different values in succession. If you do not specify `times`, the last value will be latched. Otherwise, each value will be returned the specified number of times.
 
 ```ts
-const spy = vi.fn()
+const mock = when(vi.fn()).calledWith('hello').thenReturn('hi', 'sup?')
 
-when(spy).calledWith('hello').thenReturn('hi', 'sup?')
-
-expect(spy('hello')).toEqual('hi')
-expect(spy('hello')).toEqual('sup?')
-expect(spy('hello')).toEqual('sup?')
+expect(mock('hello')).toEqual('hi')
+expect(mock('hello')).toEqual('sup?')
+expect(mock('hello')).toEqual('sup?')
 ```
 
-### `.thenResolve(value: TReturn)`
+### `.thenResolve(value: TReturn) -> Mock<TFunc>`
 
 When the stubbing is satisfied, resolve a `Promise` with `value`
 
 ```ts
-const spy = vi.fn()
-
-when(spy).calledWith('hello').thenResolve('world')
+const mock = when(vi.fn()).calledWith('hello').thenResolve('world')
 
-expect(await spy('hello')).toEqual('world')
+await expect(mock('hello')).resolves.toEqual('world')
 ```
 
 To only resolve a value once, use the `times` option.
@@ -333,36 +321,32 @@ To only resolve a value once, use the `times` option.
 ```ts
 import { when } from 'vitest-when'
 
-const spy = vi.fn()
-
-when(spy, { times: 1 }).calledWith('hello').thenResolve('world')
+const mock = when(vi.fn(), { times: 1 })
+  .calledWith('hello')
+  .thenResolve('world')
 
-expect(await spy('hello')).toEqual('world')
-expect(spy('hello')).toEqual(undefined)
+await expect(mock('hello')).resolves.toEqual('world')
+expect(mock('hello')).toEqual(undefined)
 ```
 
 You may pass several values to `thenResolve` to resolve different values in succession. If you do not specify `times`, the last value will be latched. Otherwise, each value will be resolved the specified number of times.
 
 ```ts
-const spy = vi.fn()
-
-when(spy).calledWith('hello').thenResolve('hi', 'sup?')
+const mock = when(vi.fn()).calledWith('hello').thenResolve('hi', 'sup?')
 
-expect(await spy('hello')).toEqual('hi')
-expect(await spy('hello')).toEqual('sup?')
-expect(await spy('hello')).toEqual('sup?')
+await expect(mock('hello')).resolves.toEqual('hi')
+await expect(mock('hello')).resolves.toEqual('sup?')
+await expect(mock('hello')).resolves.toEqual('sup?')
 ```
 
-### `.thenThrow(error: unknown)`
+### `.thenThrow(error: unknown) -> Mock<TFunc>`
 
 When the stubbing is satisfied, throw `error`.
 
 ```ts
-const spy = vi.fn()
+const mock = when(vi.fn()).calledWith('hello').thenThrow(new Error('oh no'))
 
-when(spy).calledWith('hello').thenThrow(new Error('oh no'))
-
-expect(() => spy('hello')).toThrow('oh no')
+expect(() => mock('hello')).toThrow('oh no')
 ```
 
 To only throw an error only once, use the `times` option.
@@ -370,38 +354,34 @@ To only throw an error only once, use the `times` option.
 ```ts
 import { when } from 'vitest-when'
 
-const spy = vi.fn()
-
-when(spy, { times: 1 }).calledWith('hello').thenThrow(new Error('oh no'))
+const mock = when(vi.fn(), { times: 1 })
+  .calledWith('hello')
+  .thenThrow(new Error('oh no'))
 
-expect(() => spy('hello')).toThrow('oh no')
-expect(spy('hello')).toEqual(undefined)
+expect(() => mock('hello')).toThrow('oh no')
+expect(mock('hello')).toEqual(undefined)
 ```
 
 You may pass several values to `thenThrow` to throw different errors in succession. If you do not specify `times`, the last value will be latched. Otherwise, each error will be thrown the specified number of times.
 
 ```ts
-const spy = vi.fn()
-
-when(spy)
+const mock = when(vi.fn())
   .calledWith('hello')
   .thenThrow(new Error('oh no'), new Error('this is bad'))
 
-expect(() => spy('hello')).toThrow('oh no')
-expect(() => spy('hello')).toThrow('this is bad')
-expect(() => spy('hello')).toThrow('this is bad')
+expect(() => mock('hello')).toThrow('oh no')
+expect(() => mock('hello')).toThrow('this is bad')
+expect(() => mock('hello')).toThrow('this is bad')
 ```
 
-### `.thenReject(error: unknown)`
+### `.thenReject(error: unknown) -> Mock<TFunc>`
 
 When the stubbing is satisfied, reject a `Promise` with `error`.
 
 ```ts
-const spy = vi.fn()
-
-when(spy).calledWith('hello').thenReject(new Error('oh no'))
+const mock = when(vi.fn()).calledWith('hello').thenReject(new Error('oh no'))
 
-await expect(spy('hello')).rejects.toThrow('oh no')
+await expect(mock('hello')).rejects.toThrow('oh no')
 ```
 
 To only throw an error only once, use the `times` option.
@@ -409,44 +389,41 @@ To only throw an error only once, use the `times` option.
 ```ts
 import { times, when } from 'vitest-when'
 
-const spy = vi.fn()
-
-when(spy, { times: 1 }).calledWith('hello').thenReject(new Error('oh no'))
+const mock = when(vi.fn(), { times: 1 })
+  .calledWith('hello')
+  .thenReject(new Error('oh no'))
 
-await expect(spy('hello')).rejects.toThrow('oh no')
-expect(spy('hello')).toEqual(undefined)
+await expect(mock('hello')).rejects.toThrow('oh no')
+expect(mock('hello')).toEqual(undefined)
 ```
 
 You may pass several values to `thenReject` to throw different errors in succession. If you do not specify `times`, the last value will be latched. Otherwise, each rejection will be triggered the specified number of times.
 
 ```ts
-const spy = vi.fn()
-
-when(spy)
+const mock = when(vi.fn())
   .calledWith('hello')
   .thenReject(new Error('oh no'), new Error('this is bad'))
 
-await expect(spy('hello')).rejects.toThrow('oh no')
-await expect(spy('hello')).rejects.toThrow('this is bad')
-await expect(spy('hello')).rejects.toThrow('this is bad')
+await expect(mock('hello')).rejects.toThrow('oh no')
+await expect(mock('hello')).rejects.toThrow('this is bad')
+await expect(mock('hello')).rejects.toThrow('this is bad')
 ```
 
-### `.thenDo(callback: (...args: TArgs) => TReturn)`
+### `.thenDo(callback: (...args: TArgs) => TReturn) -> Mock<TFunc>`
 
 When the stubbing is satisfied, run `callback` to trigger a side-effect and return its result (if any). `thenDo` is a relatively powerful tool for stubbing complex behaviors, so if you find yourself using `thenDo` often, consider refactoring your code to use more simple interactions! Your future self will thank you.
 
 ```ts
-const spy = vi.fn()
 let called = false
 
-when(spy)
+const mock = when(vi.fn())
   .calledWith('hello')
   .thenDo(() => {
     called = true
     return 'world'
   })
 
-expect(spy('hello')).toEqual('world')
+expect(mock('hello')).toEqual('world')
 expect(called).toEqual(true)
 ```
 
@@ -455,33 +432,29 @@ To only run the callback once, use the `times` option.
 ```ts
 import { times, when } from 'vitest-when'
 
-const spy = vi.fn()
-
-when(spy, { times: 1 })
+const mock = when(vi.fn(), { times: 1 })
   .calledWith('hello')
   .thenDo(() => 'world')
 
-expect(spy('hello')).toEqual('world')
-expect(spy('hello')).toEqual(undefined)
+expect(mock('hello')).toEqual('world')
+expect(mock('hello')).toEqual(undefined)
 ```
 
 You may pass several callbacks to `thenDo` to trigger different side-effects in succession. If you do not specify `times`, the last callback will be latched. Otherwise, each callback will be triggered the specified number of times.
 
 ```ts
-const spy = vi.fn()
-
-when(spy)
+const mock = when(vi.fn())
   .calledWith('hello')
   .thenDo(
     () => 'world',
     () => 'solar system',
   )
 
-expect(spy('hello')).toEqual('world')
-expect(spy('hello')).toEqual('solar system')
+expect(mock('hello')).toEqual('world')
+expect(mock('hello')).toEqual('solar system')
 ```
 
-### `debug(spy: TFunc, options?: DebugOptions): DebugResult`
+### `debug(mock: TFunc, options?: DebugOptions): DebugResult`
 
 Logs and returns information about a mock's stubbing and usage. Useful if a test with mocks is failing and you can't figure out why.
 
diff --git a/src/behaviors.ts b/src/behaviors.ts
index c93512d..7b4c5b9 100644
--- a/src/behaviors.ts
+++ b/src/behaviors.ts
@@ -3,8 +3,8 @@ import { equals } from '@vitest/expect'
 import type {
   AnyCallable,
   AnyFunction,
-  ExtractParameters,
-  ExtractReturnType,
+  ParametersOf,
+  ReturnTypeOf,
   WithMatchers,
 } from './types.ts'
 
@@ -14,17 +14,17 @@ export interface WhenOptions {
 
 export interface BehaviorStack<TFunc extends AnyCallable> {
   use: (
-    args: ExtractParameters<TFunc>,
-  ) => BehaviorEntry<ExtractParameters<TFunc>> | undefined
+    args: ParametersOf<TFunc>,
+  ) => BehaviorEntry<ParametersOf<TFunc>> | undefined
 
-  getAll: () => readonly BehaviorEntry<ExtractParameters<TFunc>>[]
+  getAll: () => readonly BehaviorEntry<ParametersOf<TFunc>>[]
 
-  getUnmatchedCalls: () => readonly ExtractParameters<TFunc>[]
+  getUnmatchedCalls: () => readonly ParametersOf<TFunc>[]
 
   bindArgs: (
-    args: WithMatchers<ExtractParameters<TFunc>>,
+    args: WithMatchers<ParametersOf<TFunc>>,
     options: WhenOptions,
-  ) => BoundBehaviorStack<ExtractReturnType<TFunc>>
+  ) => BoundBehaviorStack<ReturnTypeOf<TFunc>>
 }
 
 export interface BoundBehaviorStack<TReturn> {
@@ -65,8 +65,8 @@ export interface BehaviorOptions<TValue> {
 export const createBehaviorStack = <
   TFunc extends AnyCallable,
 >(): BehaviorStack<TFunc> => {
-  const behaviors: BehaviorEntry<ExtractParameters<TFunc>>[] = []
-  const unmatchedCalls: ExtractParameters<TFunc>[] = []
+  const behaviors: BehaviorEntry<ParametersOf<TFunc>>[] = []
+  const unmatchedCalls: ParametersOf<TFunc>[] = []
 
   return {
     getAll: () => behaviors,
diff --git a/src/debug.ts b/src/debug.ts
index 986a1bd..364411d 100644
--- a/src/debug.ts
+++ b/src/debug.ts
@@ -3,9 +3,9 @@ import {
   plugins as prettyFormatPlugins,
 } from 'pretty-format'
 
-import { type Behavior, BehaviorType } from './behaviors'
-import { getBehaviorStack, validateSpy } from './stubs'
-import type { AnyCallable, MockInstance } from './types'
+import { type Behavior, BehaviorType } from './behaviors.ts'
+import { getBehaviorStack } from './stubs.ts'
+import type { AnyCallable, Mock } from './types.ts'
 
 export interface DebugResult {
   name: string
@@ -21,12 +21,11 @@ export interface Stubbing {
 }
 
 export const getDebug = <TFunc extends AnyCallable>(
-  spy: TFunc | MockInstance<TFunc>,
+  mock: Mock<TFunc>,
 ): DebugResult => {
-  const target = validateSpy<TFunc>(spy)
-  const name = target.getMockName()
-  const behaviors = getBehaviorStack(target)
-  const unmatchedCalls = behaviors?.getUnmatchedCalls() ?? target.mock.calls
+  const name = mock.getMockName()
+  const behaviors = getBehaviorStack(mock)
+  const unmatchedCalls = behaviors?.getUnmatchedCalls() ?? mock.mock.calls
   const stubbings =
     behaviors?.getAll().map((entry) => ({
       args: entry.args,
diff --git a/src/fallback-implementation.ts b/src/fallback-implementation.ts
index 9e915be..6073263 100644
--- a/src/fallback-implementation.ts
+++ b/src/fallback-implementation.ts
@@ -1,17 +1,18 @@
-import type { AnyCallable, MockInstance } from './types.ts'
+import type { AnyCallable, 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: MockInstance<TFunc>,
-): TFunc | undefined => {
+  mock: Mock<TFunc>,
+): AsFunction<TFunc> | undefined => {
   return (
-    mock.getMockImplementation() ?? getTinyspyInternals(mock)?.getOriginal()
+    (mock.getMockImplementation() as AsFunction<TFunc> | undefined) ??
+    getTinyspyInternals(mock)?.getOriginal()
   )
 }
 
 /** Internal state from Tinyspy, where a mock's default implementation is stored. */
 interface TinyspyInternals<TFunc extends AnyCallable> {
-  getOriginal: () => TFunc | undefined
+  getOriginal: () => AsFunction<TFunc> | undefined
 }
 
 /**
@@ -24,7 +25,7 @@ interface TinyspyInternals<TFunc extends AnyCallable> {
  * which is stored on a Symbol key in the mock object.
  */
 const getTinyspyInternals = <TFunc extends AnyCallable>(
-  mock: MockInstance<TFunc>,
+  mock: Mock<TFunc>,
 ): TinyspyInternals<TFunc> | undefined => {
   const maybeTinyspy = mock as unknown as Record<PropertyKey, unknown>
 
diff --git a/src/stubs.ts b/src/stubs.ts
index fe92e1d..9352545 100644
--- a/src/stubs.ts
+++ b/src/stubs.ts
@@ -5,37 +5,31 @@ import {
 } from './behaviors.ts'
 import { NotAMockFunctionError } from './errors.ts'
 import { getFallbackImplementation } from './fallback-implementation.ts'
-import type {
-  AnyCallable,
-  AnyFunction,
-  ExtractParameters,
-  MockInstance,
-} from './types.ts'
+import type { AnyCallable, Mock, ParametersOf } from './types.ts'
 
-const BEHAVIORS_KEY = Symbol('behaviors')
+const BEHAVIORS_KEY = Symbol.for('vitest-when:behaviors')
 
 interface WhenStubImplementation<TFunc extends AnyCallable> {
-  (...args: ExtractParameters<TFunc>): unknown
+  (...args: ParametersOf<TFunc>): unknown
   [BEHAVIORS_KEY]: BehaviorStack<TFunc>
 }
 
-export const configureStub = <TFunc extends AnyCallable>(
-  maybeSpy: unknown,
+export const configureMock = <TFunc extends AnyCallable>(
+  mock: Mock<TFunc>,
 ): BehaviorStack<TFunc> => {
-  const spy = validateSpy<TFunc>(maybeSpy)
-  const existingBehaviors = getBehaviorStack(spy)
+  const existingBehaviorStack = getBehaviorStack(mock)
 
-  if (existingBehaviors) {
-    return existingBehaviors
+  if (existingBehaviorStack) {
+    return existingBehaviorStack
   }
 
-  const behaviors = createBehaviorStack<TFunc>()
-  const fallbackImplementation = getFallbackImplementation(spy)
+  const behaviorStack = createBehaviorStack<TFunc>()
+  const fallbackImplementation = getFallbackImplementation(mock)
 
-  const implementation = (...args: ExtractParameters<TFunc>) => {
-    const behavior = behaviors.use(args)?.behavior ?? {
+  const implementation = (...args: ParametersOf<TFunc>) => {
+    const behavior = behaviorStack.use(args)?.behavior ?? {
       type: BehaviorType.DO,
-      callback: fallbackImplementation as AnyFunction | undefined,
+      callback: fallbackImplementation,
     }
 
     switch (behavior.type) {
@@ -62,35 +56,31 @@ export const configureStub = <TFunc extends AnyCallable>(
     }
   }
 
-  spy.mockImplementation(
-    Object.assign(implementation as TFunc, { [BEHAVIORS_KEY]: behaviors }),
+  mock.mockImplementation(
+    Object.assign(implementation, { [BEHAVIORS_KEY]: behaviorStack }),
   )
 
-  return behaviors
+  return behaviorStack
 }
 
-export const validateSpy = <TFunc extends AnyCallable>(
-  maybeSpy: unknown,
-): MockInstance<TFunc> => {
+export const validateMock = <TFunc extends AnyCallable>(
+  maybeMock: TFunc | Mock<TFunc>,
+): Mock<TFunc> => {
   if (
-    typeof maybeSpy === 'function' &&
-    'mockImplementation' in maybeSpy &&
-    typeof maybeSpy.mockImplementation === 'function' &&
-    'getMockImplementation' in maybeSpy &&
-    typeof maybeSpy.getMockImplementation === 'function' &&
-    'getMockName' in maybeSpy &&
-    typeof maybeSpy.getMockName === 'function'
+    typeof maybeMock === 'function' &&
+    'mockImplementation' in maybeMock &&
+    typeof maybeMock.mockImplementation === 'function'
   ) {
-    return maybeSpy as unknown as MockInstance<TFunc>
+    return maybeMock
   }
 
-  throw new NotAMockFunctionError(maybeSpy)
+  throw new NotAMockFunctionError(maybeMock)
 }
 
 export const getBehaviorStack = <TFunc extends AnyCallable>(
-  spy: MockInstance<TFunc>,
+  mock: Mock<TFunc>,
 ): BehaviorStack<TFunc> | undefined => {
-  const existingImplementation = spy.getMockImplementation() as
+  const existingImplementation = mock.getMockImplementation() as
     | WhenStubImplementation<TFunc>
     | TFunc
     | undefined
diff --git a/src/types.ts b/src/types.ts
index bf48046..74bbfb1 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -1,5 +1,6 @@
 /** Common type definitions. */
 import type { AsymmetricMatcher } from '@vitest/expect'
+import type { MockedClass, MockedFunction } from 'vitest'
 
 /** Any function. */
 export type AnyFunction = (...args: never[]) => unknown
@@ -11,37 +12,34 @@ export type AnyConstructor = new (...args: never[]) => unknown
 export type AnyCallable = AnyFunction | AnyConstructor
 
 /** Extract parameters from either a function or constructor. */
-export type ExtractParameters<T> = T extends new (...args: infer P) => unknown
+export type ParametersOf<TFunc extends AnyCallable> = TFunc extends new (
+  ...args: infer P
+) => unknown
   ? P
-  : T extends (...args: infer P) => unknown
+  : TFunc extends (...args: infer P) => unknown
     ? P
     : never
 
 /** Extract return type from either a function or constructor */
-export type ExtractReturnType<T> = T extends new (...args: never[]) => infer R
+export type ReturnTypeOf<TFunc extends AnyCallable> = TFunc extends new (
+  ...args: never[]
+) => infer R
   ? R
-  : T extends (...args: never[]) => infer R
+  : TFunc extends (...args: never[]) => infer R
     ? R
     : never
 
+export type AsFunction<TFunc extends AnyCallable> = (
+  ...args: ParametersOf<TFunc>
+) => ReturnTypeOf<TFunc>
+
 /** Accept a value or an AsymmetricMatcher in an arguments array */
 export type WithMatchers<T extends unknown[]> = {
   [K in keyof T]: AsymmetricMatcher<unknown> | T[K]
 }
 
-/**
- * Minimally typed version of Vitest's `MockInstance`.
- *
- * Used to ensure backwards compatibility
- * with older versions of Vitest.
- */
-export interface MockInstance<TFunc extends AnyCallable = AnyCallable> {
-  getMockName(): string
-  getMockImplementation(): TFunc | undefined
-  mockImplementation: (impl: TFunc) => this
-  mock: MockContext<TFunc>
-}
-
-export interface MockContext<TFunc extends AnyCallable> {
-  calls: ExtractParameters<TFunc>[]
-}
+export type Mock<TFunc extends AnyCallable> = TFunc extends AnyFunction
+  ? MockedFunction<TFunc>
+  : TFunc extends AnyConstructor
+    ? MockedClass<TFunc>
+    : never
diff --git a/src/vitest-when.ts b/src/vitest-when.ts
index fdc1fef..499f97f 100644
--- a/src/vitest-when.ts
+++ b/src/vitest-when.ts
@@ -1,11 +1,12 @@
 import type { WhenOptions } from './behaviors.ts'
 import { type DebugResult, getDebug } from './debug.ts'
-import { configureStub } from './stubs.ts'
+import { configureMock, validateMock } from './stubs.ts'
 import type {
   AnyCallable,
-  ExtractParameters,
-  ExtractReturnType,
-  MockInstance,
+  AsFunction,
+  Mock,
+  ParametersOf,
+  ReturnTypeOf,
   WithMatchers,
 } from './types.ts'
 
@@ -14,35 +15,51 @@ export type { DebugResult, Stubbing } from './debug.ts'
 export * from './errors.ts'
 
 export interface StubWrapper<TFunc extends AnyCallable> {
-  calledWith<TArgs extends ExtractParameters<TFunc>>(
+  calledWith<TArgs extends ParametersOf<TFunc>>(
     ...args: WithMatchers<TArgs>
-  ): Stub<TArgs, ExtractReturnType<TFunc>>
+  ): Stub<TFunc>
 }
 
-export interface Stub<TArgs extends unknown[], TReturn> {
-  thenReturn: (...values: TReturn[]) => void
-  thenResolve: (...values: Awaited<TReturn>[]) => void
-  thenThrow: (...errors: unknown[]) => void
-  thenReject: (...errors: unknown[]) => void
-  thenDo: (...callbacks: ((...args: TArgs) => TReturn)[]) => void
+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 const when = <TFunc extends AnyCallable>(
-  spy: TFunc | MockInstance<TFunc>,
+  mock: TFunc | Mock<TFunc>,
   options: WhenOptions = {},
 ): StubWrapper<TFunc> => {
-  const behaviorStack = configureStub(spy)
+  const validatedMock = validateMock(mock)
+  const behaviorStack = configureMock(validatedMock)
 
   return {
     calledWith: (...args) => {
       const behaviors = behaviorStack.bindArgs(args, options)
 
       return {
-        thenReturn: (...values) => behaviors.addReturn(values),
-        thenResolve: (...values) => behaviors.addResolve(values),
-        thenThrow: (...errors) => behaviors.addThrow(errors),
-        thenReject: (...errors) => behaviors.addReject(errors),
-        thenDo: (...callbacks) => behaviors.addDo(callbacks),
+        thenReturn: (...values) => {
+          behaviors.addReturn(values)
+          return validatedMock
+        },
+        thenResolve: (...values) => {
+          behaviors.addResolve(values)
+          return validatedMock
+        },
+        thenThrow: (...errors) => {
+          behaviors.addThrow(errors)
+          return validatedMock
+        },
+        thenReject: (...errors) => {
+          behaviors.addReject(errors)
+          return validatedMock
+        },
+        thenDo: (...callbacks) => {
+          behaviors.addDo(callbacks)
+          return validatedMock
+        },
       }
     },
   }
@@ -53,11 +70,12 @@ export interface DebugOptions {
 }
 
 export const debug = <TFunc extends AnyCallable>(
-  spy: TFunc | MockInstance<TFunc>,
+  mock: TFunc | Mock<TFunc>,
   options: DebugOptions = {},
 ): DebugResult => {
   const log = options.log ?? true
-  const result = getDebug(spy)
+  const validatedMock = validateMock(mock)
+  const result = getDebug(validatedMock)
 
   if (log) {
     console.debug(result.description)
diff --git a/test/typing.test-d.ts b/test/typing.test-d.ts
index 5471b3e..527c8e0 100644
--- a/test/typing.test-d.ts
+++ b/test/typing.test-d.ts
@@ -1,68 +1,109 @@
 /* eslint-disable
+  @typescript-eslint/require-await,
   @typescript-eslint/no-explicit-any,
   @typescript-eslint/restrict-template-expressions
 */
 
-import { assertType, describe, expect, it, vi } from 'vitest'
+import { describe, expect, expectTypeOf, it, vi } from 'vitest'
 
 import * as subject from '../src/vitest-when.ts'
 
 describe('vitest-when type signatures', () => {
   it('should handle an anonymous mock', () => {
-    const spy = vi.fn()
-    const stub = subject.when(spy).calledWith(1, 2, 3)
+    const result = subject.when(vi.fn()).calledWith(1, 2, 3).thenReturn(4)
 
-    assertType<subject.Stub<[number, number, number], any>>(stub)
+    expectTypeOf(result).parameters.toEqualTypeOf<any[]>()
+    expectTypeOf(result).returns.toEqualTypeOf<any>()
+    expectTypeOf(result.mock.calls).toEqualTypeOf<any[][]>()
   })
 
   it('should handle an untyped function', () => {
-    const stub = subject.when(untyped).calledWith(1)
+    const result = subject.when(untyped).calledWith(1).thenReturn('hello')
 
-    stub.thenReturn('hello')
+    expectTypeOf(result.mock.calls).toEqualTypeOf<any[][]>()
+    expectTypeOf(result).parameters.toEqualTypeOf<any[]>()
+    expectTypeOf(result).returns.toEqualTypeOf<any>()
+  })
+
+  it('should handle a simple function', () => {
+    const result = subject.when(simple).calledWith(1).thenReturn('hello')
 
-    assertType<subject.Stub<[number], any>>(stub)
+    expectTypeOf(result.mock.calls).toEqualTypeOf<[number][]>()
+    expectTypeOf(result).parameters.toEqualTypeOf<[number]>()
+    expectTypeOf(result).returns.toEqualTypeOf<string>()
   })
 
-  it('should handle an spied function', () => {
-    const target = { simple }
-    const spy = vi.spyOn(target, 'simple')
-    const stub = subject.when(spy).calledWith(1)
+  it('returns mock type for then resolve', () => {
+    const result = subject.when(simpleAsync).calledWith(1).thenResolve('hello')
+
+    expectTypeOf(result.mock.calls).toEqualTypeOf<[number][]>()
+    expectTypeOf(result).parameters.toEqualTypeOf<[number]>()
+    expectTypeOf(result).returns.toEqualTypeOf<Promise<string>>()
+  })
 
-    stub.thenReturn('hello')
+  it('returns mock type for then throw', () => {
+    const result = subject.when(simple).calledWith(1).thenThrow('oh no')
 
-    assertType<subject.Stub<[number], any>>(stub)
+    expectTypeOf(result.mock.calls).toEqualTypeOf<[number][]>()
+    expectTypeOf(result).parameters.toEqualTypeOf<[number]>()
+    expectTypeOf(result).returns.toEqualTypeOf<string>()
   })
 
-  it('should handle a simple function', () => {
-    const stub = subject.when(simple).calledWith(1)
+  it('returns mock type for then reject', () => {
+    const result = subject.when(simpleAsync).calledWith(1).thenReject('oh no')
+
+    expectTypeOf(result.mock.calls).toEqualTypeOf<[number][]>()
+    expectTypeOf(result).parameters.toEqualTypeOf<[number]>()
+    expectTypeOf(result).returns.toEqualTypeOf<Promise<string>>()
+  })
 
-    stub.thenReturn('hello')
+  it('returns mock type for then do', () => {
+    const result = subject
+      .when(simple)
+      .calledWith(1)
+      .thenDo(() => 'hello')
 
-    assertType<subject.Stub<[1], string>>(stub)
+    expectTypeOf(result.mock.calls).toEqualTypeOf<[number][]>()
+    expectTypeOf(result).parameters.toEqualTypeOf<[number]>()
+    expectTypeOf(result).returns.toEqualTypeOf<string>()
   })
 
-  it('should handle a generic function', () => {
-    const stub = subject.when(generic).calledWith(1)
+  it('should handle an spied function', () => {
+    const target = { simple }
+    vi.spyOn(target, 'simple')
 
-    stub.thenReturn('hello')
+    const result = subject.when(target.simple).calledWith(1).thenReturn('hello')
 
-    assertType<subject.Stub<[number], string>>(stub)
+    expectTypeOf(result.mock.calls).toEqualTypeOf<[number][]>()
+    expectTypeOf(result).parameters.toEqualTypeOf<[number]>()
+    expectTypeOf(result).returns.toEqualTypeOf<string>()
   })
 
-  it('should handle an overloaded function using its last overload', () => {
-    const stub = subject.when(overloaded).calledWith(1)
-
-    stub.thenReturn('hello')
+  it('should handle a generic function', () => {
+    const result = subject.when(generic).calledWith(1).thenReturn('hello')
 
-    assertType<subject.Stub<[1], string>>(stub)
+    expectTypeOf(result.mock.calls).toEqualTypeOf<[unknown][]>()
+    expectTypeOf(result).parameters.toEqualTypeOf<[unknown]>()
+    expectTypeOf(result).returns.toEqualTypeOf<string>()
   })
 
-  it('should handle an overloaded function using an explicit type', () => {
-    const stub = subject.when<() => boolean>(overloaded).calledWith()
+  it('should handle an overloaded function using its last overload', () => {
+    const result = subject.when(overloaded).calledWith(1).thenReturn('hello')
 
-    stub.thenReturn(true)
+    expectTypeOf(result.mock.calls).toEqualTypeOf<[number][]>()
+    expectTypeOf(result).parameters.toEqualTypeOf<[number]>()
+    expectTypeOf(result).returns.toEqualTypeOf<string>()
+  })
 
-    assertType<subject.Stub<[], boolean>>(stub)
+  it('should handle an overloaded function using an explicit type', () => {
+    const result = subject
+      .when<() => boolean>(overloaded)
+      .calledWith()
+      .thenReturn(true)
+
+    expectTypeOf(result.mock.calls).toEqualTypeOf<[][]>()
+    expectTypeOf(result).parameters.toEqualTypeOf<[]>()
+    expectTypeOf(result).returns.toEqualTypeOf<boolean>()
   })
 
   it('should reject invalid usage of a simple function', () => {
@@ -100,7 +141,13 @@ describe('vitest-when type signatures', () => {
       }
     }
 
-    subject.when(TestClass).calledWith(42)
+    const result = subject
+      .when(TestClass)
+      .calledWith(42)
+      .thenReturn({} as TestClass)
+
+    expectTypeOf(result.mock.instances).toEqualTypeOf<TestClass[]>()
+    expectTypeOf(result).constructorParameters.toEqualTypeOf<[number]>()
 
     // @ts-expect-error: args wrong type
     subject.when(TestClass).calledWith('42')
@@ -115,6 +162,10 @@ function simple(input: number): string {
   throw new Error(`simple(${input})`)
 }
 
+async function simpleAsync(input: number): Promise<string> {
+  throw new Error(`simpleAsync(${input})`)
+}
+
 function complex(input: { a: number; b: string }): string {
   throw new Error(`simple({ a: ${input.a}, b: ${input.b} })`)
 }
diff --git a/test/vitest-when.test.ts b/test/vitest-when.test.ts
index 35c6524..063df06 100644
--- a/test/vitest-when.test.ts
+++ b/test/vitest-when.test.ts
@@ -27,9 +27,7 @@ describe('vitest-when', () => {
   })
 
   it('should return undefined by default', () => {
-    const spy = vi.fn()
-
-    subject.when(spy).calledWith(1, 2, 3).thenReturn(4)
+    const spy = subject.when(vi.fn()).calledWith(1, 2, 3).thenReturn(4)
 
     expect(spy()).toEqual(undefined)
     expect(spy(1)).toEqual(undefined)
@@ -39,26 +37,23 @@ describe('vitest-when', () => {
   })
 
   it('should return a value', () => {
-    const spy = vi.fn()
-
-    subject.when(spy).calledWith(1, 2, 3).thenReturn(4)
+    const spy = subject.when(vi.fn()).calledWith(1, 2, 3).thenReturn(4)
 
     expect(spy(1, 2, 3)).toEqual(4)
     expect(spy(1, 2, 3)).toEqual(4)
   })
 
   it('should return undefined if passed nothing', () => {
-    const spy = vi.fn()
-
-    subject.when(spy).calledWith(1, 2, 3).thenReturn()
+    const spy = subject.when(vi.fn()).calledWith(1, 2, 3).thenReturn()
 
     expect(spy(1, 2, 3)).toEqual(undefined)
   })
 
   it('should fall back to original mock implementation', () => {
-    const spy = vi.fn().mockReturnValue(100)
-
-    subject.when(spy).calledWith(1, 2, 3).thenReturn(4)
+    const spy = subject
+      .when(vi.fn().mockReturnValue(100))
+      .calledWith(1, 2, 3)
+      .thenReturn(4)
 
     expect(spy(1, 2, 3)).toEqual(4)
     expect(spy()).toEqual(100)
@@ -76,9 +71,10 @@ describe('vitest-when', () => {
   })
 
   it('should return a number of times', () => {
-    const spy = vi.fn()
-
-    subject.when(spy, { times: 2 }).calledWith(1, 2, 3).thenReturn(4)
+    const spy = subject
+      .when(vi.fn(), { times: 2 })
+      .calledWith(1, 2, 3)
+      .thenReturn(4)
 
     expect(spy(1, 2, 3)).toEqual(4)
     expect(spy(1, 2, 3)).toEqual(4)
@@ -86,18 +82,17 @@ describe('vitest-when', () => {
   })
 
   it('should be resettable', () => {
-    const spy = vi.fn()
-
-    subject.when(spy).calledWith(1, 2, 3).thenReturn(4)
+    const spy = subject.when(vi.fn()).calledWith(1, 2, 3).thenReturn(4)
     vi.resetAllMocks()
 
     expect(spy(1, 2, 3)).toEqual(undefined)
   })
 
   it('should throw an error', () => {
-    const spy = vi.fn()
-
-    subject.when(spy).calledWith(1, 2, 3).thenThrow(new Error('oh no'))
+    const spy = subject
+      .when(vi.fn())
+      .calledWith(1, 2, 3)
+      .thenThrow(new Error('oh no'))
 
     expect(() => {
       spy(1, 2, 3)
@@ -105,34 +100,29 @@ describe('vitest-when', () => {
   })
 
   it('should resolve a Promise', async () => {
-    const spy = vi.fn()
-
-    subject.when(spy).calledWith(1, 2, 3).thenResolve(4)
+    const spy = subject.when(vi.fn()).calledWith(1, 2, 3).thenResolve(4)
 
     await expect(spy(1, 2, 3)).resolves.toEqual(4)
   })
 
   it('should resolve undefined if passed nothing', async () => {
-    const spy = vi.fn()
-
-    subject.when(spy).calledWith(1, 2, 3).thenResolve()
+    const spy = subject.when(vi.fn()).calledWith(1, 2, 3).thenResolve()
 
     await expect(spy(1, 2, 3)).resolves.toEqual(undefined)
   })
 
   it('should reject a Promise', async () => {
-    const spy = vi.fn()
-
-    subject.when(spy).calledWith(1, 2, 3).thenReject(new Error('oh no'))
+    const spy = subject
+      .when(vi.fn())
+      .calledWith(1, 2, 3)
+      .thenReject(new Error('oh no'))
 
     await expect(spy(1, 2, 3)).rejects.toThrow('oh no')
   })
 
   it('should do a callback', () => {
-    const spy = vi.fn()
     const callback = vi.fn(() => 4)
-
-    subject.when(spy).calledWith(1, 2, 3).thenDo(callback)
+    const spy = subject.when(vi.fn()).calledWith(1, 2, 3).thenDo(callback)
 
     expect(spy(1, 2, 3)).toEqual(4)
     expect(callback).toHaveBeenCalledWith(1, 2, 3)
@@ -140,9 +130,7 @@ describe('vitest-when', () => {
   })
 
   it('should return multiple values', () => {
-    const spy = vi.fn()
-
-    subject.when(spy).calledWith(1, 2, 3).thenReturn(4, 5, 6)
+    const spy = subject.when(vi.fn()).calledWith(1, 2, 3).thenReturn(4, 5, 6)
 
     expect(spy(1, 2, 3)).toEqual(4)
     expect(spy(1, 2, 3)).toEqual(5)
@@ -151,9 +139,7 @@ describe('vitest-when', () => {
   })
 
   it('should resolve multiple values', async () => {
-    const spy = vi.fn()
-
-    subject.when(spy).calledWith(1, 2, 3).thenResolve(4, 5, 6)
+    const spy = subject.when(vi.fn()).calledWith(1, 2, 3).thenResolve(4, 5, 6)
 
     await expect(spy(1, 2, 3)).resolves.toEqual(4)
     await expect(spy(1, 2, 3)).resolves.toEqual(5)
@@ -162,10 +148,8 @@ describe('vitest-when', () => {
   })
 
   it('should reject multiple errors', async () => {
-    const spy = vi.fn()
-
-    subject
-      .when(spy)
+    const spy = subject
+      .when(vi.fn())
       .calledWith(1, 2, 3)
       .thenReject(new Error('4'), new Error('5'), new Error('6'))
 
@@ -176,10 +160,8 @@ describe('vitest-when', () => {
   })
 
   it('should reject a number of times', async () => {
-    const spy = vi.fn()
-
-    subject
-      .when(spy, { times: 2 })
+    const spy = subject
+      .when(vi.fn(), { times: 2 })
       .calledWith(1, 2, 3)
       .thenReject(new Error('4'))
 
@@ -189,32 +171,31 @@ describe('vitest-when', () => {
   })
 
   it('should throw multiple errors', () => {
-    const spy = vi.fn()
-
-    subject
-      .when(spy)
+    const spy = subject
+      .when(vi.fn())
       .calledWith(1, 2, 3)
       .thenThrow(new Error('4'), new Error('5'), new Error('6'))
 
     expect(() => {
       spy(1, 2, 3)
     }).toThrow('4')
+
     expect(() => {
       spy(1, 2, 3)
     }).toThrow('5')
+
     expect(() => {
       spy(1, 2, 3)
     }).toThrow('6')
+
     expect(() => {
       spy(1, 2, 3)
     }).toThrow('6')
   })
 
   it('should call multiple callbacks', () => {
-    const spy = vi.fn()
-
-    subject
-      .when(spy)
+    const spy = subject
+      .when(vi.fn())
       .calledWith(1, 2, 3)
       .thenDo(
         () => 4,
@@ -229,9 +210,7 @@ describe('vitest-when', () => {
   })
 
   it('should allow multiple different stubs', () => {
-    const spy = vi.fn()
-
-    subject.when(spy).calledWith(1, 2, 3).thenReturn(4)
+    const spy = subject.when(vi.fn()).calledWith(1, 2, 3).thenReturn(4)
     subject.when(spy).calledWith(4, 5, 6).thenReturn(7)
 
     expect(spy(1, 2, 3)).toEqual(4)
@@ -239,19 +218,15 @@ describe('vitest-when', () => {
   })
 
   it('should use the latest stub', () => {
-    const spy = vi.fn()
-
-    subject.when(spy).calledWith(1, 2, 3).thenReturn(4)
+    const spy = subject.when(vi.fn()).calledWith(1, 2, 3).thenReturn(4)
     subject.when(spy).calledWith(1, 2, 3).thenReturn(1000)
 
     expect(spy(1, 2, 3)).toEqual(1000)
   })
 
   it('should respect asymmetric matchers', () => {
-    const spy = vi.fn()
-
-    subject
-      .when(spy)
+    const spy = subject
+      .when(vi.fn())
       .calledWith(expect.stringContaining('foo'))
       .thenReturn(1000)
 
@@ -259,18 +234,17 @@ describe('vitest-when', () => {
   })
 
   it('should respect custom asymmetric matchers', () => {
-    const spy = vi.fn()
-
-    subject.when(spy).calledWith(expect.toBeFoo()).thenReturn(1000)
+    const spy = subject
+      .when(vi.fn())
+      .calledWith(expect.toBeFoo())
+      .thenReturn(1000)
 
     expect(spy('foo')).toEqual(1000)
   })
 
   it('should deeply check object arguments', () => {
-    const spy = vi.fn()
-
-    subject
-      .when(spy)
+    const spy = subject
+      .when(vi.fn())
       .calledWith({ foo: { bar: { baz: 0 } } })
       .thenReturn(100)
 
@@ -278,9 +252,9 @@ describe('vitest-when', () => {
   })
 
   it('should not trigger unhandled rejection warnings when rejection unused', () => {
-    const spy = vi.fn()
     const error = new Error('uh uhh')
-    subject.when(spy).calledWith('/api/foo').thenReject(error)
+    subject.when(vi.fn()).calledWith('/api/foo').thenReject(error)
+
     // intentionally do not call the spy
     expect(true).toBe(true)
   })