diff --git a/CHANGELOG.md b/CHANGELOG.md index 1a758cb98c90..84cb35b41293 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,7 @@ - `[@jest/environment, jest-runtime]` Allow `jest.requireActual` and `jest.requireMock` to take a type argument ([#13253](https://github.com/facebook/jest/pull/13253)) - `[@jest/environment]` Allow `jest.mock` and `jest.doMock` to take a type argument ([#13254](https://github.com/facebook/jest/pull/13254)) - `[@jest/fake-timers]` Add `jest.now()` to return the current fake clock time ([#13244](https://github.com/facebook/jest/pull/13244), [13246](https://github.com/facebook/jest/pull/13246)) +- `[@jest/mock]` Add `withImplementation` method for temporarily overriding a mock. ### Fixes diff --git a/docs/MockFunctionAPI.md b/docs/MockFunctionAPI.md index 8e55737147d8..43a95bb3192e 100644 --- a/docs/MockFunctionAPI.md +++ b/docs/MockFunctionAPI.md @@ -478,6 +478,43 @@ test('async test', async () => { }); ``` +### `mockFn.withImplementation(fn, callback)` + +Accepts a function which should be temporarily used as the implementation of the mock while the callback is being executed. + +```js +test('test', () => { + const mock = jest.fn(() => 'outside callback'); + + mock.withImplementation( + () => 'inside callback', + () => { + mock(); // 'inside callback' + }, + ); + + mock(); // 'outside callback' +}); +``` + +`mockFn.withImplementation` can be used regardless of whether or not the callback is asynchronous (returns a `thenable`). If the callback is asynchronous a promise will be returned. Awaiting the promise will await the callback and reset the implementation. + +```js +test('async test', async () => { + const mock = jest.fn(() => 'outside callback'); + + // We await this call since the callback is async + await mock.withImplementation( + () => 'inside callback', + async () => { + mock(); // 'inside callback' + }, + ); + + mock(); // 'outside callback' +}); +``` + ## TypeScript Usage :::tip diff --git a/packages/jest-mock/README.md b/packages/jest-mock/README.md index 5bc75093f8fa..41588a07b281 100644 --- a/packages/jest-mock/README.md +++ b/packages/jest-mock/README.md @@ -98,3 +98,9 @@ In case both `.mockImplementationOnce()` / `.mockImplementation()` and `.mockRet - if the last call is `.mockReturnValueOnce()` or `.mockReturnValue()`, use the specific return value or default return value. If specific return values are used up or no default return value is set, fall back to try `.mockImplementation()`; - if the last call is `.mockImplementationOnce()` or `.mockImplementation()`, run the specific implementation and return the result or run default implementation and return the result. + +##### `.withImplementation(function, callback)` + +Temporarily overrides the default mock implementation within the callback, then restores it's previous implementation. + +If the callback is async or returns a `thenable`, `withImplementation` will return a promise. Awaiting the promise will await the callback and reset the implementation. diff --git a/packages/jest-mock/__typetests__/mock-functions.test.ts b/packages/jest-mock/__typetests__/mock-functions.test.ts index 0c747c2642e0..3d72cc4d37d0 100644 --- a/packages/jest-mock/__typetests__/mock-functions.test.ts +++ b/packages/jest-mock/__typetests__/mock-functions.test.ts @@ -250,6 +250,13 @@ expectType Promise>>( ); expectError(fn(() => Promise.resolve('')).mockRejectedValueOnce()); +expectType(mockFn.withImplementation(mockFnImpl, () => {})); +expectType>( + mockFn.withImplementation(mockFnImpl, async () => {}), +); + +expectError(mockFn.withImplementation(mockFnImpl)); + // jest.spyOn() const spiedArray = ['a', 'b']; diff --git a/packages/jest-mock/package.json b/packages/jest-mock/package.json index ba0ca81eab42..bd24a84b3037 100644 --- a/packages/jest-mock/package.json +++ b/packages/jest-mock/package.json @@ -18,7 +18,8 @@ }, "dependencies": { "@jest/types": "workspace:^", - "@types/node": "*" + "@types/node": "*", + "jest-util": "workspace:^" }, "devDependencies": { "@tsd/typescript": "~4.8.2", diff --git a/packages/jest-mock/src/__tests__/index.test.ts b/packages/jest-mock/src/__tests__/index.test.ts index 6b9f716d7ef0..27f7ba18369c 100644 --- a/packages/jest-mock/src/__tests__/index.test.ts +++ b/packages/jest-mock/src/__tests__/index.test.ts @@ -8,6 +8,7 @@ /* eslint-disable local/ban-types-eventually, local/prefer-rest-params-eventually */ +import * as util from 'util'; import vm, {Context} from 'vm'; import {ModuleMocker, fn, mocked, spyOn} from '../'; @@ -1073,6 +1074,62 @@ describe('moduleMocker', () => { }); }); + describe('withImplementation', () => { + it('sets an implementation which is available within the callback', () => { + const mock1 = jest.fn(); + const mock2 = jest.fn(); + + const Module = jest.fn(() => ({someFn: mock1})); + const testFn = function () { + const m = new Module(); + m.someFn(); + }; + + Module.withImplementation( + () => ({someFn: mock2}), + () => { + testFn(); + expect(mock2).toHaveBeenCalled(); + expect(mock1).not.toHaveBeenCalled(); + }, + ); + + testFn(); + expect(mock1).toHaveBeenCalled(); + + expect.assertions(3); + }); + + it('returns a promise if the provided callback is asynchronous', async () => { + const mock1 = jest.fn(); + const mock2 = jest.fn(); + + const Module = jest.fn(() => ({someFn: mock1})); + const testFn = function () { + const m = new Module(); + m.someFn(); + }; + + const promise = Module.withImplementation( + () => ({someFn: mock2}), + async () => { + testFn(); + expect(mock2).toHaveBeenCalled(); + expect(mock1).not.toHaveBeenCalled(); + }, + ); + + expect(util.types.isPromise(promise)).toBe(true); + + await promise; + + testFn(); + expect(mock1).toHaveBeenCalled(); + + expect.assertions(4); + }); + }); + test('mockReturnValue does not override mockImplementationOnce', () => { const mockFn = jest .fn() diff --git a/packages/jest-mock/src/index.ts b/packages/jest-mock/src/index.ts index 69319a05c583..b727d18e749f 100644 --- a/packages/jest-mock/src/index.ts +++ b/packages/jest-mock/src/index.ts @@ -7,6 +7,8 @@ /* eslint-disable local/ban-types-eventually, local/prefer-rest-params-eventually */ +import {isPromise} from 'jest-util'; + export type MockMetadataType = | 'object' | 'array' @@ -135,6 +137,8 @@ export interface MockInstance { mockRestore(): void; mockImplementation(fn: T): this; mockImplementationOnce(fn: T): this; + withImplementation(fn: T, callback: () => Promise): Promise; + withImplementation(fn: T, callback: () => void): void; mockName(name: string): this; mockReturnThis(): this; mockReturnValue(value: ReturnType): this; @@ -768,6 +772,34 @@ export class ModuleMocker { return f; }; + f.withImplementation = withImplementation.bind(this); + + function withImplementation(fn: T, callback: () => void): void; + function withImplementation( + fn: T, + callback: () => Promise, + ): Promise; + function withImplementation( + this: ModuleMocker, + fn: T, + callback: (() => void) | (() => Promise), + ): void | Promise { + // Remember previous mock implementation, then set new one + const mockConfig = this._ensureMockConfig(f); + const previousImplementation = mockConfig.mockImpl; + mockConfig.mockImpl = fn; + + const returnedValue = callback(); + + if (isPromise(returnedValue)) { + return returnedValue.then(() => { + mockConfig.mockImpl = previousImplementation; + }); + } else { + mockConfig.mockImpl = previousImplementation; + } + } + f.mockImplementation = (fn: UnknownFunction) => { // next function call will use mock implementation return value const mockConfig = this._ensureMockConfig(f); diff --git a/packages/jest-mock/tsconfig.json b/packages/jest-mock/tsconfig.json index b69d4caaeea9..cf4cceccce14 100644 --- a/packages/jest-mock/tsconfig.json +++ b/packages/jest-mock/tsconfig.json @@ -6,5 +6,5 @@ }, "include": ["./src/**/*"], "exclude": ["./**/__tests__/**/*"], - "references": [{"path": "../jest-types"}] + "references": [{"path": "../jest-types"}, {"path": "../jest-util"}] } diff --git a/yarn.lock b/yarn.lock index 071ce29e6769..ecb380ee4efe 100644 --- a/yarn.lock +++ b/yarn.lock @@ -12578,6 +12578,7 @@ __metadata: "@jest/types": "workspace:^" "@tsd/typescript": ~4.8.2 "@types/node": "*" + jest-util: "workspace:^" tsd-lite: ^0.6.0 languageName: unknown linkType: soft