Skip to content

Commit

Permalink
fix(spy): fix mockImplementation for function overload and unions (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
hi-ogawa authored Jul 22, 2024
1 parent f68453f commit 7a75bd4
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 4 deletions.
10 changes: 6 additions & 4 deletions packages/spy/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,8 @@ export interface MockContext<T extends Procedure> {
}

type Procedure = (...args: any[]) => any
// pick a single function type from function overloads, unions, etc...
type NormalizedPrecedure<T extends Procedure> = (...args: Parameters<T>) => ReturnType<T>

type Methods<T> = keyof {
[K in keyof T as T[K] extends Procedure ? K : never]: T[K];
Expand Down Expand Up @@ -204,22 +206,22 @@ export interface MockInstance<T extends Procedure = Procedure> {
*
* If mock was created with `vi.spyOn`, it will return `undefined` unless a custom implementation was provided.
*/
getMockImplementation(): T | undefined
getMockImplementation(): NormalizedPrecedure<T> | undefined
/**
* Accepts a function that will be used as an implementation of the mock.
* @example
* const increment = vi.fn().mockImplementation(count => count + 1);
* expect(increment(3)).toBe(4);
*/
mockImplementation(fn: T): this
mockImplementation(fn: NormalizedPrecedure<T>): this
/**
* Accepts a function that will be used as a mock implementation during the next call. Can be chained so that multiple function calls produce different results.
* @example
* const fn = vi.fn(count => count).mockImplementationOnce(count => count + 1);
* expect(fn(3)).toBe(4);
* expect(fn(3)).toBe(3);
*/
mockImplementationOnce(fn: T): this
mockImplementationOnce(fn: NormalizedPrecedure<T>): this
/**
* Overrides the original mock implementation temporarily while the callback is being executed.
* @example
Expand All @@ -231,7 +233,7 @@ export interface MockInstance<T extends Procedure = Procedure> {
*
* myMockFn() // 'original'
*/
withImplementation<T2>(fn: T, cb: () => T2): T2 extends Promise<unknown> ? Promise<this> : this
withImplementation<T2>(fn: NormalizedPrecedure<T>, cb: () => T2): T2 extends Promise<unknown> ? Promise<this> : this

/**
* Use this if you need to return `this` context from the method without invoking actual implementation.
Expand Down
16 changes: 16 additions & 0 deletions test/core/test/vi.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -152,6 +152,22 @@ describe('testing vi utils', () => {
expectTypeOf(gSpy.mock.contexts).toEqualTypeOf<unknown[]>()
})

test('mockImplementation types', async () => {
// overload
const fs = { readFileSync() {} } as any as typeof import('node:fs')
vi.spyOn(fs, 'readFileSync').mockImplementation(() => 'str')
vi.spyOn(fs, 'readFileSync').mockImplementation(() => Buffer.from('buf'))
vi.fn(fs.readFileSync).mockImplementation(() => 'str')
vi.fn(fs.readFileSync).mockImplementation(() => Buffer.from('buf'))

// union
interface Handler {
(v: number): number
other: (v: number) => number
}
vi.fn<Handler>().mockImplementation(v => v + 1)
})

test('can change config', () => {
const state = getWorkerState()
expect(state.config.hookTimeout).toBe(10000)
Expand Down

0 comments on commit 7a75bd4

Please sign in to comment.