Add .thenCallback() stubbing

Author: sirlancelot

src/behaviors.ts

@@ -1,6 +1,12 @@
 import { equals } from '@vitest/expect'
 
-import type { AnyMockable, ParametersOf, WithMatchers } from './types.ts'
+import { concatImpliedCallback, invokeCallbackFor } from './callback.ts'
+import type {
+  AnyFunction,
+  AnyMockable,
+  ParametersOf,
+  WithMatchers,
+} from './types.ts'
 
 export interface WhenOptions {
   ignoreExtraArgs?: boolean
@@ -12,20 +18,20 @@ export interface PlanOptions extends WhenOptions {
 }
 
 export type StubbingPlan =
+  | 'thenCallback'
   | 'thenDo'
   | 'thenReject'
   | 'thenResolve'
   | 'thenReturn'
   | 'thenThrow'
 
 export interface UseAction {
-  callCount: number
   plan: StubbingPlan
-  values: unknown[]
+  value: unknown
 }
 
 export interface BehaviorStack<TFunc extends AnyMockable> {
-  use: (args: ParametersOf<TFunc>, fallbackValue: unknown) => UseAction
+  use: (args: ParametersOf<TFunc>, fallbackValue?: AnyFunction) => UseAction
 
   getAll: () => readonly BehaviorEntry<ParametersOf<TFunc>>[]
 
@@ -65,31 +71,32 @@ export const createBehaviorStack = <
 
       if (!behavior) {
         unmatchedCalls.push(args)
-        return { callCount: 0, plan: 'thenDo', values: [fallbackValue] }
+        return { plan: 'thenDo', value: fallbackValue }
       }
 
+      const value = getCurrentStubbedValue(behavior)
+
       behavior.calls.push(args)
-      return {
-        callCount: behavior.callCount++,
-        plan: behavior.options.plan,
-        values: behavior.values,
-      }
+      behavior.callCount++
+
+      invokeCallbackFor(behavior, args)
+
+      return { plan: behavior.options.plan, value }
     },
 
-    addStubbing: (args, values, options) => {
-      behaviors.unshift({
-        args,
-        callCount: 0,
-        calls: [],
-        options,
-        values,
-      })
+    addStubbing: (rawArgs, values, options) => {
+      const args =
+        options.plan === 'thenCallback'
+          ? concatImpliedCallback(rawArgs)
+          : rawArgs
+
+      behaviors.unshift({ args, callCount: 0, calls: [], options, values })
     },
   }
 }
 
-const behaviorMatches = <TArgs extends unknown[]>(actualArguments: TArgs) => {
-  return (behavior: BehaviorEntry<TArgs>): boolean => {
+const behaviorMatches = (actualArguments: unknown[]) => {
+  return (behavior: BehaviorEntry<unknown[]>): boolean => {
     const { callCount, options } = behavior
     const { times } = options
 
@@ -108,3 +115,11 @@ const behaviorMatches = <TArgs extends unknown[]>(actualArguments: TArgs) => {
     })
   }
 }
+
+const getCurrentStubbedValue = (
+  behavior: BehaviorEntry<unknown[]>,
+): unknown => {
+  const { callCount, values } = behavior
+  const hasMoreValues = callCount < values.length
+  return hasMoreValues ? values[callCount] : values.at(-1)
+}

src/callback.ts

@@ -0,0 +1,97 @@
+import { expect } from 'vitest'
+
+import type { BehaviorEntry } from './behaviors.ts'
+import type { AnyFunction } from './types.ts'
+
+const CALLBACK_KEY = Symbol.for('vitest-when:callback-args')
+
+export interface WithCallbackMarker {
+  [CALLBACK_KEY]: unknown[]
+}
+
+declare module 'vitest' {
+  interface ExpectStatic {
+    callback: (...args: unknown[]) => WithCallbackMarker
+  }
+}
+
+/**
+ * Creates a matcher that matches a callback function. Passing no arguments
+ * creates an empty callback matcher, in which case the arguments will be picked
+ * up from the call to `.thenCallback()`.
+ */
+export function createCallbackMatcher(...args: unknown[]) {
+  const matcher = expect.any(Function) as WithCallbackMarker
+  matcher[CALLBACK_KEY] = args
+  return matcher
+}
+
+expect.callback = createCallbackMatcher
+
+export function isCallback(value: unknown): value is WithCallbackMarker {
+  return (
+    typeof value === 'object' &&
+    value !== null &&
+    CALLBACK_KEY in value &&
+    value[CALLBACK_KEY] !== undefined
+  )
+}
+
+/**
+ * Append an empty callback matcher if one isn't already present.
+ */
+export function concatImpliedCallback<TArgs extends unknown[]>(
+  args: TArgs,
+): TArgs {
+  const callbackArguments = args.filter(isCallback)
+
+  // The user didn't provide a callback matcher, so we'll add one
+  if (callbackArguments.length === 0)
+    return [...args, expect.callback()] as TArgs
+
+  // The user provided one or more callback matchers, and at least one of them
+  // is the special empty one.
+  for (const callbackArgument of callbackArguments) {
+    const stubbedArgs = callbackArgument[CALLBACK_KEY]
+    if (stubbedArgs.length === 0) return args
+  }
+
+  // The user provided one or more callback matchers, but none of them are the
+  // special empty one, so we'll add one.
+  return [...args, expect.callback()] as TArgs
+}
+
+/**
+ * Invokes any callbacks that were stubbed for the given behavior.
+ */
+export const invokeCallbackFor = (
+  behavior: BehaviorEntry<unknown[]>,
+  actualArguments: unknown[],
+) => {
+  for (const [index, expectedArgument] of behavior.args.entries()) {
+    if (!isCallback(expectedArgument)) continue
+
+    // Callback here is guaranteed to be a function due to the matcher
+    // validation in `createCallbackMatcher()`, so we can cast it as such.
+    const callback = actualArguments[index] as AnyFunction
+
+    // If this is the empty callback matcher, then we'll use the values passed
+    // to `.thenCallback()` as the arguments.
+    const callbackArguments =
+      expectedArgument[CALLBACK_KEY].length === 0
+        ? behavior.values
+        : expectedArgument[CALLBACK_KEY]
+
+    // Defer the callback to the next tick to ensure correct async behavior as
+    // well as protect the core from any userland errors.
+    setImmediate(callback, ...callbackArguments)
+  }
+}
+
+/**
+ * Returns a Promise that resolves on the next tick. When testing callbacks, all
+ * invoked callbacks are deferred until the next tick, so you must `await
+ * nextTick()` before asserting on the callback arguments.
+ */
+export const nextTick = () =>
+  new Promise<void>((resolve) => setImmediate(resolve))

src/debug.ts

@@ -8,11 +8,6 @@ import type { StubbingPlan } from './behaviors.ts'
 import { getBehaviorStack } from './stubs.ts'
 import type { MockInstance } from './types.ts'
 
-interface Behavior {
-  type: StubbingPlan
-  values: readonly unknown[]
-}
-
 export interface DebugResult {
   name: string
   description: string

src/stubs.ts

@@ -29,15 +29,11 @@ export const configureMock = <TFunc extends AnyMockable>(
   const behaviorStack = createBehaviorStack<TFunc>()
   const fallbackImplementation = getFallbackImplementation(mock)
 
-  function implementation(this: ThisType<TFunc>, ...args: ParametersOf<TFunc>) {
-    const { callCount, plan, values } = behaviorStack.use(
-      args,
-      fallbackImplementation,
-    )
-    const hasMoreValues = callCount < values.length
-
-    // eslint-disable-next-line @typescript-eslint/no-unsafe-call
-    const value: unknown = hasMoreValues ? values[callCount] : values.at(-1)
+  function implementation(
+    this: ThisType<TFunc>,
+    ...args: ParametersOf<TFunc>
+  ): unknown {
+    const { plan, value } = behaviorStack.use(args, fallbackImplementation)
 
     switch (plan) {
       case 'thenReturn': {
@@ -54,11 +50,15 @@ export const configureMock = <TFunc extends AnyMockable>(
         return Promise.reject(value)
       }
       case 'thenDo': {
-        // eslint-disable-next-line @typescript-eslint/no-unsafe-return
         return typeof value === 'function'
           ? value.call(this, ...args)
           : undefined
       }
+      case 'thenCallback': {
+        // Callback is called during `behaviorStack.use()` so there's nothing to
+        // do here.
+        return undefined
+      }
     }
   }
 

src/vitest-when.ts

@@ -13,6 +13,10 @@ import type {
 } from './types.ts'
 
 export type { WhenOptions } from './behaviors.ts'
+export {
+  createCallbackMatcher as expectCallback,
+  nextTick,
+} from './callback.ts'
 export type { DebugResult, Stubbing } from './debug.ts'
 export * from './errors.ts'
 
@@ -28,6 +32,7 @@ export interface Stub<TFunc extends AnyMockable> {
   thenThrow: (...errors: unknown[]) => Mock<TFunc>
   thenReject: (...errors: unknown[]) => Mock<TFunc>
   thenDo: (...callbacks: AsFunction<TFunc>[]) => Mock<TFunc>
+  thenCallback: (...args: unknown[]) => Mock<TFunc>
 }
 
 export const when = <TFunc extends AnyMockable>(
@@ -76,6 +81,13 @@ export const when = <TFunc extends AnyMockable>(
           })
           return result
         },
+        thenCallback: (...callbackArguments) => {
+          behaviorStack.addStubbing(args, callbackArguments, {
+            ...options,
+            plan: 'thenCallback',
+          })
+          return result
+        },
       }
     },
   }

test/callback.test.ts

@@ -0,0 +1,64 @@
+import { describe, expect, it, vi } from 'vitest'
+
+import * as subject from '../src/vitest-when.ts'
+
+describe('vitest-when callback', () => {
+  it('invokes an implied callback', async () => {
+    const next = vi.fn()
+    const spy = subject
+      .when(vi.fn())
+      .calledWith('hello', 'world')
+      .thenCallback(undefined, 'okay')
+
+    expect(spy('hello', 'world', next)).toBeUndefined()
+
+    await subject.nextTick()
+    expect(next).toHaveBeenCalledExactlyOnceWith(undefined, 'okay')
+  })
+
+  it('invokes a specified callback', async () => {
+    const next = vi.fn()
+    const spy = subject
+      .when(vi.fn())
+      .calledWith('hello', expect.callback(), 'world')
+      .thenCallback('okay')
+
+    spy('hello', next, 'world')
+
+    await subject.nextTick()
+    expect(next).toHaveBeenCalledExactlyOnceWith('okay')
+  })
+
+  it('invokes multiple specified callbacks', async () => {
+    const next = vi.fn()
+    const spy = subject
+      .when(vi.fn())
+      .calledWith(
+        expect.callback(undefined, 'onStart'),
+        expect.callback(undefined, 'onEnd'),
+      )
+      .thenReturn('okay')
+
+    expect(spy(next, next)).toBe('okay')
+
+    await subject.nextTick()
+    expect(next).toHaveBeenCalledTimes(2)
+    expect(next).toHaveBeenNthCalledWith(1, undefined, 'onStart')
+    expect(next).toHaveBeenNthCalledWith(2, undefined, 'onEnd')
+  })
+
+  it('appends implied callback even when specified callbacks are provided', async () => {
+    const next = vi.fn()
+    const spy = subject
+      .when(vi.fn())
+      .calledWith(expect.callback('first'))
+      .thenCallback('second')
+
+    spy(next, next)
+
+    await subject.nextTick()
+    expect(next).toHaveBeenCalledTimes(2)
+    expect(next).toHaveBeenNthCalledWith(1, 'first')
+    expect(next).toHaveBeenNthCalledWith(2, 'second')
+  })
+})

test/debug.test.ts

@@ -90,9 +90,7 @@ describe('vitest-when debug', () => {
         },
       ],
       unmatchedCalls: [[1234]],
-      description: expect.stringMatching(
-        /1 call:\n\s+- \(1234\)/,
-      ) as string,
+      description: expect.stringMatching(/1 call:\n\s+- \(1234\)/) as string,
     } satisfies subject.DebugResult)
   })