Adjust when type to support ignoreExtraArgs

Author: sirlancelot

src/behaviors.ts

@@ -9,8 +9,8 @@ import type {
   WithMatchers,
 } from './types.ts'
 
-export interface WhenOptions {
-  ignoreExtraArgs?: boolean
+export interface WhenOptions<TIgnore extends boolean = boolean> {
+  ignoreExtraArgs?: TIgnore
   times?: number
 }
 

src/vitest-when.ts

@@ -22,6 +22,12 @@ export interface StubWrapper<TFunc extends AnyMockable> {
   ): Stub<TFunc>
 }
 
+export interface StubWrapperFlexible<TFunc extends AnyMockable> {
+  calledWith<TArgs extends ParametersOf<TFunc>>(
+    ...args: Partial<WithMatchers<TArgs>>
+  ): Stub<TFunc>
+}
+
 export interface Stub<TFunc extends AnyMockable> {
   thenReturn: (...values: ReturnTypeOf<TFunc>[]) => Mock<TFunc>
   thenResolve: (...values: Awaited<ReturnTypeOf<TFunc>>[]) => Mock<TFunc>
@@ -30,10 +36,18 @@ export interface Stub<TFunc extends AnyMockable> {
   thenDo: (...callbacks: AsFunction<TFunc>[]) => Mock<TFunc>
 }
 
-export const when = <TFunc extends AnyMockable>(
+export function when<TFunc extends AnyMockable>(
+  mock: TFunc | MockInstance<TFunc>,
+  options: WhenOptions<true>,
+): StubWrapperFlexible<NormalizeMockable<TFunc>>
+export function when<TFunc extends AnyMockable>(
+  mock: TFunc | MockInstance<TFunc>,
+  options?: WhenOptions<false>,
+): StubWrapper<NormalizeMockable<TFunc>>
+export function when<TFunc extends AnyMockable>(
   mock: TFunc | MockInstance<TFunc>,
   options: WhenOptions = {},
-): StubWrapper<NormalizeMockable<TFunc>> => {
+): StubWrapper<NormalizeMockable<TFunc>> {
   const validatedMock = validateMock(mock)
   const behaviorStack = configureMock(validatedMock)
   const result = asMock(validatedMock)

test/fixtures.ts

@@ -31,6 +31,10 @@ export function overloaded(input?: number): string | boolean {
   throw new Error(`overloaded(${input})`)
 }
 
+export function extraArguments(first: number, second: number): string {
+  throw new Error(`extraArguments(${first}, ${second})`)
+}
+
 export class SimpleClass {
   constructor(input: number) {
     throw new Error(`SimpleClass(${input})`)

test/typing.test-d.ts

@@ -16,6 +16,7 @@ import {
 import * as subject from '../src/vitest-when.ts'
 import {
   complex,
+  extraArguments,
   generic,
   overloaded,
   simple,
@@ -45,6 +46,17 @@ describe('vitest-when type signatures', () => {
     >()
   })
 
+  it('should handle fewer than required arguments', () => {
+    const result = subject
+      .when(extraArguments, { ignoreExtraArgs: true })
+      .calledWith(1)
+      .thenReturn('hello')
+
+    expectTypeOf(result).toEqualTypeOf<
+      MockedFunction<(input: number, second: number) => string>
+    >()
+  })
+
   it('returns mock type for then resolve', () => {
     const result = subject.when(simpleAsync).calledWith(1).thenResolve('hello')