Skip to content

Commit

Permalink
fix(spy): don't allow Promise in mockImplementation (#4859)
Browse files Browse the repository at this point in the history
  • Loading branch information
sheremet-va authored Jan 3, 2024
1 parent 1fee63f commit 072e02b
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 6 deletions.
12 changes: 6 additions & 6 deletions packages/spy/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -154,15 +154,15 @@ export interface MockInstance<TArgs extends any[] = any[], TReturns = any> {
* const increment = vi.fn().mockImplementation(count => count + 1);
* expect(increment(3)).toBe(4);
*/
mockImplementation(fn: ((...args: TArgs) => TReturns) | (() => Promise<TReturns>)): this
mockImplementation(fn: ((...args: TArgs) => TReturns)): 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: ((...args: TArgs) => TReturns) | (() => Promise<TReturns>)): this
mockImplementationOnce(fn: ((...args: TArgs) => TReturns)): this
/**
* Overrides the original mock implementation temporarily while the callback is being executed.
* @example
Expand Down Expand Up @@ -479,16 +479,16 @@ function enhanceSpy<TArgs extends any[], TReturns>(
stub.mockReturnValueOnce = (val: TReturns) => stub.mockImplementationOnce(() => val)

stub.mockResolvedValue = (val: Awaited<TReturns>) =>
stub.mockImplementation(() => Promise.resolve(val as TReturns))
stub.mockImplementation(() => Promise.resolve(val as TReturns) as any)

stub.mockResolvedValueOnce = (val: Awaited<TReturns>) =>
stub.mockImplementationOnce(() => Promise.resolve(val as TReturns))
stub.mockImplementationOnce(() => Promise.resolve(val as TReturns) as any)

stub.mockRejectedValue = (val: unknown) =>
stub.mockImplementation(() => Promise.reject(val))
stub.mockImplementation(() => Promise.reject(val) as any)

stub.mockRejectedValueOnce = (val: unknown) =>
stub.mockImplementationOnce(() => Promise.reject(val))
stub.mockImplementationOnce(() => Promise.reject(val) as any)

Object.defineProperty(stub, 'mock', {
get: () => mockContext,
Expand Down
21 changes: 21 additions & 0 deletions test/core/test/jest-mock.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,27 @@ describe('jest mock compat layer', () => {
expect(mock2.getMockImplementation()).toBeUndefined()
})

it('implementation types allow only function returned types', () => {
function fn() {
return 1
}

function asyncFn() {
return Promise.resolve(1)
}

const mock1 = vi.fn(fn)
const mock2 = vi.fn(asyncFn)

mock1.mockImplementation(() => 2)
// @ts-expect-error promise is not allowed
mock1.mockImplementation(() => Promise.resolve(2))

// @ts-expect-error non-promise is not allowed
mock2.mockImplementation(() => 2)
mock2.mockImplementation(() => Promise.resolve(2))
})

it('implementation sync fn', () => {
const originalFn = function () {
return 'original'
Expand Down

0 comments on commit 072e02b

Please sign in to comment.