Skip to content

Commit e0167ea

Browse files
committed
fix: improve spying types
1 parent c57511b commit e0167ea

File tree

4 files changed

+175
-18
lines changed

4 files changed

+175
-18
lines changed

packages/spy/src/index.ts

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -24,10 +24,10 @@ const MOCK_RESTORE = new Set<() => void>()
2424
// Jest keeps the state in a separate WeakMap which is good for memory,
2525
// but it makes the state slower to access and return different values
2626
// if you stored it before calling `mockClear` where it will be recreated
27-
const REGISTERED_MOCKS = new Set<Mock>()
28-
const MOCK_CONFIGS = new WeakMap<Mock, MockConfig>()
27+
const REGISTERED_MOCKS = new Set<Mock<Procedure | Constructable>>()
28+
const MOCK_CONFIGS = new WeakMap<Mock<Procedure | Constructable>, MockConfig>()
2929

30-
export function createMockInstance(options: MockInstanceOption = {}): Mock {
30+
export function createMockInstance(options: MockInstanceOption = {}): Mock<Procedure | Constructable> {
3131
const {
3232
originalImplementation,
3333
restore,
@@ -225,7 +225,7 @@ export function spyOn<T extends object, K extends keyof T>(
225225
object: T,
226226
key: K,
227227
accessor?: 'get' | 'set',
228-
): Mock {
228+
): Mock<Procedure | Constructable> {
229229
assert(
230230
object != null,
231231
'The vi.spyOn() function could not find an object to spy upon. The first argument must be defined.',
@@ -498,10 +498,11 @@ function createMock(
498498
return returnValue
499499
}) as Mock,
500500
}
501+
const mock = namedObject[name] as Mock<Procedure | Constructable>
501502
if (original) {
502-
copyOriginalStaticProperties(namedObject[name], original)
503+
copyOriginalStaticProperties(mock, original)
503504
}
504-
return namedObject[name]
505+
return mock
505506
}
506507

507508
function registerCalls(args: unknown[], state: MockContext, prototypeState?: MockContext) {
@@ -536,7 +537,7 @@ function registerContext(context: MockProcedureContext<Procedure>, state: MockCo
536537
return [contextIndex, contextPrototypeIndex] as const
537538
}
538539

539-
function copyOriginalStaticProperties(mock: Mock, original: Procedure | Constructable) {
540+
function copyOriginalStaticProperties(mock: Mock<Procedure | Constructable>, original: Procedure | Constructable) {
540541
const { properties, descriptors } = getAllProperties(original)
541542

542543
for (const key of properties) {
@@ -625,6 +626,7 @@ export function resetAllMocks(): void {
625626
}
626627

627628
export type {
629+
Constructable,
628630
MaybeMocked,
629631
MaybeMockedConstructor,
630632
MaybeMockedDeep,
@@ -654,4 +656,5 @@ export type {
654656
PartiallyMockedFunction,
655657
PartiallyMockedFunctionDeep,
656658
PartialMock,
659+
Procedure,
657660
} from './types'

packages/spy/src/types.ts

Lines changed: 30 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,7 @@ export interface MockContext<T extends Procedure | Constructable = Procedure> {
7676
* This is an array containing all instances that were instantiated when mock was called with a `new` keyword. Note that this is an actual context (`this`) of the function, not a return value.
7777
* @see https://vitest.dev/api/mock#mock-instances
7878
*/
79-
instances: MockReturnType<T>[]
79+
instances: MockProcedureContext<T>[]
8080
/**
8181
* An array of `this` values that were used during each call to the mock function.
8282
* @see https://vitest.dev/api/mock#mock-contexts
@@ -284,7 +284,8 @@ export interface MockInstance<T extends Procedure | Constructable = Procedure> e
284284
*
285285
* myMockFn() // 'original'
286286
*/
287-
withImplementation<T2>(fn: NormalizedProcedure<T>, cb: () => T2): T2 extends Promise<unknown> ? Promise<this> : this
287+
withImplementation(fn: NormalizedProcedure<T>, cb: () => Promise<unknown>): Promise<this>
288+
withImplementation(fn: NormalizedProcedure<T>, cb: () => unknown): this
288289

289290
/**
290291
* Use this if you need to return the `this` context from the method without invoking the actual implementation.
@@ -359,12 +360,29 @@ export interface MockInstance<T extends Procedure | Constructable = Procedure> e
359360
}
360361
/* eslint-enable ts/method-signature-style */
361362

362-
export interface Mock<T extends Procedure | Constructable = Procedure> extends MockInstance<T> {
363-
new (...args: MockParameters<T>): T extends Constructable ? InstanceType<T> : MockReturnType<T>
364-
(...args: MockParameters<T>): MockReturnType<T>
363+
export type Mock<T extends Procedure | Constructable = Procedure> = MockInstance<T> & {
365364
/** @internal */
366365
_isMockFunction: true
367-
}
366+
} & (
367+
T extends Constructable
368+
? (
369+
T extends Procedure
370+
// supports both `new Class()` and `Class()`
371+
? {
372+
new (...args: ConstructorParameters<T>): InstanceType<T>
373+
(...args: Parameters<T>): ReturnType<T>
374+
}
375+
// supports only `new Class()`
376+
: {
377+
new (...args: ConstructorParameters<T>): InstanceType<T>
378+
}
379+
)
380+
// any function can be called with the new keyword
381+
: {
382+
new (...args: MockParameters<T>): MockReturnType<T>
383+
(...args: MockParameters<T>): MockReturnType<T>
384+
}
385+
) & { [P in keyof T]: T[P] }
368386

369387
type PartialMaybePromise<T> = T extends Promise<Awaited<T>>
370388
? Promise<Partial<Awaited<T>>>
@@ -381,12 +399,13 @@ type PartialResultFunction<T> = T extends Constructable
381399
? (...args: Parameters<T>) => PartialMaybePromise<ReturnType<T>>
382400
: T
383401

384-
export interface PartialMock<T extends Procedure | Constructable = Procedure>
385-
extends Mock<
386-
PartialResultFunction<T extends Mock
402+
export type PartialMock<T extends Procedure | Constructable = Procedure> = Mock<
403+
PartialResultFunction<
404+
T extends Mock
387405
? NonNullable<ReturnType<T['getMockImplementation']>>
388-
: T>
389-
> {}
406+
: T
407+
>
408+
>
390409

391410
export type MaybeMockedConstructor<T> = T extends Constructable
392411
? Mock<T>
Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
import type { Mock, MockResult, MockSettledResult, Procedure } from '@vitest/spy'
2+
import { expectTypeOf, test } from 'vitest'
3+
4+
test('spy.mock when implementation is a class', () => {
5+
class Klass {
6+
constructor(_a: string, _b?: number) {
7+
// ...
8+
}
9+
10+
static getType() {
11+
return 'Klass'
12+
}
13+
}
14+
15+
const Mock = vi.fn(Klass)
16+
17+
expectTypeOf(Mock.mock.calls).toEqualTypeOf<[a: string, b?: number][]>()
18+
expectTypeOf(Mock.mock.results).toEqualTypeOf<MockResult<void>[]>()
19+
expectTypeOf(Mock.mock.contexts).toEqualTypeOf<Klass[]>()
20+
expectTypeOf(Mock.mock.instances).toEqualTypeOf<Klass[]>()
21+
expectTypeOf(Mock.mock.invocationCallOrder).toEqualTypeOf<number[]>()
22+
expectTypeOf(Mock.mock.settledResults).toEqualTypeOf<MockSettledResult<void>[]>()
23+
expectTypeOf(Mock.mock.lastCall).toEqualTypeOf<[a: string, b?: number] | undefined>()
24+
25+
// static properties are defined
26+
expectTypeOf(Mock.getType).toBeFunction()
27+
expectTypeOf(Mock.getType).returns.toBeString()
28+
29+
expectTypeOf(Mock).constructorParameters.toEqualTypeOf<[a: string, b?: number]>()
30+
expectTypeOf(Mock).instance.toEqualTypeOf<Klass>()
31+
})
32+
33+
test('spy.mock when implementation is a class-like function', () => {
34+
function Klass(this: typeof Klass, _a: string, _b?: number) {
35+
// ...
36+
}
37+
38+
const Mock = vi.fn(Klass)
39+
40+
expectTypeOf(Mock.mock.calls).toEqualTypeOf<[a: string, b?: number][]>()
41+
expectTypeOf(Mock.mock.results).toEqualTypeOf<MockResult<void>[]>()
42+
expectTypeOf(Mock.mock.contexts).toEqualTypeOf<typeof Klass[]>()
43+
expectTypeOf(Mock.mock.instances).toEqualTypeOf<typeof Klass[]>()
44+
expectTypeOf(Mock.mock.invocationCallOrder).toEqualTypeOf<number[]>()
45+
expectTypeOf(Mock.mock.settledResults).toEqualTypeOf<MockSettledResult<void>[]>()
46+
expectTypeOf(Mock.mock.lastCall).toEqualTypeOf<[a: string, b?: number] | undefined>()
47+
48+
expectTypeOf(Mock).constructorParameters.toEqualTypeOf<[a: string, b?: number]>()
49+
})
50+
51+
test('spy.mock when implementation is a normal function', () => {
52+
function FN(_a: string, _b?: number) {
53+
return 42
54+
}
55+
56+
const Mock = vi.fn(FN)
57+
58+
expectTypeOf(Mock.mock.calls).toEqualTypeOf<[a: string, b?: number][]>()
59+
expectTypeOf(Mock.mock.results).toEqualTypeOf<MockResult<number>[]>()
60+
expectTypeOf(Mock.mock.contexts).toEqualTypeOf<unknown[]>()
61+
expectTypeOf(Mock.mock.instances).toEqualTypeOf<unknown[]>()
62+
expectTypeOf(Mock.mock.invocationCallOrder).toEqualTypeOf<number[]>()
63+
expectTypeOf(Mock.mock.settledResults).toEqualTypeOf<MockSettledResult<number>[]>()
64+
expectTypeOf(Mock.mock.lastCall).toEqualTypeOf<[a: string, b?: number] | undefined>()
65+
66+
expectTypeOf(Mock).constructorParameters.toEqualTypeOf<[a: string, b?: number]>()
67+
})
68+
69+
test('cann call a function mock with and without new', () => {
70+
const Mock = vi.fn(function fn(this: any) {
71+
this.test = true
72+
})
73+
74+
const _mockClass = new Mock()
75+
const _mockFn = Mock()
76+
})
77+
78+
test('cannot call class mock without new', () => {
79+
const Mock = vi.fn(class {})
80+
81+
const _mockClass = new Mock()
82+
// @ts-expect-error value is not callable
83+
const _mockFn = Mock()
84+
})
85+
86+
test('spying on a function that supports new', () => {
87+
interface ReturnClass {}
88+
interface BothFnAndClass {
89+
new (): ReturnClass
90+
(): ReturnClass
91+
}
92+
93+
const Mock = vi.fn(function R() {} as BothFnAndClass)
94+
95+
// supports new
96+
const _mockClass = new Mock()
97+
// supports T()
98+
const _mockFn = Mock()
99+
})
100+
101+
test('withImplementation returns correct type', () => {
102+
const spy = vi.fn()
103+
104+
const result42 = spy.withImplementation(() => {}, () => {
105+
return 42
106+
})
107+
expectTypeOf(result42).toEqualTypeOf<Mock<Procedure>>()
108+
109+
const resultObject = spy.withImplementation(() => {}, () => {
110+
return { then: () => 42 }
111+
})
112+
expectTypeOf(resultObject).toEqualTypeOf<Mock<Procedure>>()
113+
114+
const resultVoid = spy.withImplementation(() => {}, () => {})
115+
expectTypeOf(resultVoid).toEqualTypeOf<Mock<Procedure>>()
116+
117+
const promise42 = spy.withImplementation(() => {}, async () => {
118+
return 42
119+
})
120+
expectTypeOf(promise42).toEqualTypeOf<Promise<Mock<Procedure>>>()
121+
122+
const promiseObject = spy.withImplementation(() => {}, async () => {
123+
return { hello: () => 42 }
124+
})
125+
expectTypeOf(promiseObject).toEqualTypeOf<Promise<Mock<Procedure>>>()
126+
127+
const promiseVoid = spy.withImplementation(() => {}, async () => {})
128+
expectTypeOf(promiseVoid).toEqualTypeOf<Promise<Mock<Procedure>>>()
129+
130+
const promisePromise = spy.withImplementation(() => {}, () => {
131+
return Promise.resolve()
132+
})
133+
expectTypeOf(promisePromise).toEqualTypeOf<Promise<Mock<Procedure>>>()
134+
})

test/core/test/mocking/vi-fn.test.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -652,6 +652,7 @@ describe('vi.fn() implementations', () => {
652652

653653
test('vi.fn() throws an error if new is not called on a class', () => {
654654
const Mock = vi.fn(class _Mock {})
655+
// @ts-expect-error value is not callable
655656
expect(() => Mock()).toThrowError(
656657
`Class constructor _Mock cannot be invoked without 'new'`,
657658
)

0 commit comments

Comments
 (0)