Skip to content

Commit ca041f5

Browse files
authored
fix: improve spying types (#8878)
1 parent 9e24a59 commit ca041f5

File tree

5 files changed

+186
-20
lines changed

5 files changed

+186
-20
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: 31 additions & 13 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.
@@ -356,15 +357,31 @@ export interface MockInstance<T extends Procedure | Constructable = Procedure> e
356357
* await asyncMock() // throws Error<'Async error'>
357358
*/
358359
mockRejectedValueOnce(error: unknown): this
359-
}
360-
/* eslint-enable ts/method-signature-style */
361-
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>
365360
/** @internal */
366361
_isMockFunction: true
367362
}
363+
/* eslint-enable ts/method-signature-style */
364+
365+
export type Mock<T extends Procedure | Constructable = Procedure> = MockInstance<T> & (
366+
T extends Constructable
367+
? (
368+
T extends Procedure
369+
// supports both `new Class()` and `Class()`
370+
? {
371+
new (...args: ConstructorParameters<T>): InstanceType<T>
372+
(...args: Parameters<T>): ReturnType<T>
373+
}
374+
// supports only `new Class()`
375+
: {
376+
new (...args: ConstructorParameters<T>): InstanceType<T>
377+
}
378+
)
379+
// any function can be called with the new keyword
380+
: {
381+
new (...args: MockParameters<T>): MockReturnType<T>
382+
(...args: MockParameters<T>): MockReturnType<T>
383+
}
384+
) & { [P in keyof T]: T[P] }
368385

369386
type PartialMaybePromise<T> = T extends Promise<Awaited<T>>
370387
? Promise<Partial<Awaited<T>>>
@@ -381,12 +398,13 @@ type PartialResultFunction<T> = T extends Constructable
381398
? (...args: Parameters<T>) => PartialMaybePromise<ReturnType<T>>
382399
: T
383400

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

391409
export type MaybeMockedConstructor<T> = T extends Constructable
392410
? Mock<T>

packages/vitest/src/public/index.ts

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -130,6 +130,14 @@ export type {
130130
MockedFunction,
131131
MockedObject,
132132
MockInstance,
133+
MockResult,
134+
MockResultIncomplete,
135+
MockResultReturn,
136+
MockResultThrow,
137+
MockSettledResult,
138+
MockSettledResultFulfilled,
139+
MockSettledResultIncomplete,
140+
MockSettledResultRejected,
133141
} from '@vitest/spy'
134142

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

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)