feat(@jest/mock): Add withImplementation

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
 

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

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.

packages/jest-mock/__typetests__/mock-functions.test.ts

@@ -250,6 +250,13 @@ expectType<Mock<() => Promise<string>>>(
 );
 expectError(fn(() => Promise.resolve('')).mockRejectedValueOnce());
 
+expectType<void>(mockFn.withImplementation(mockFnImpl, () => {}));
+expectType<Promise<void>>(
+  mockFn.withImplementation(mockFnImpl, async () => {}),
+);
+
+expectError(mockFn.withImplementation(mockFnImpl));
+
 // jest.spyOn()
 
 const spiedArray = ['a', '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()

packages/jest-mock/src/index.ts

@@ -135,6 +135,8 @@ export interface MockInstance<T extends FunctionLike = UnknownFunction> {
   mockRestore(): void;
   mockImplementation(fn: T): this;
   mockImplementationOnce(fn: T): this;
+  withImplementation(fn: T, callback: () => Promise<unknown>): Promise<void>;
+  withImplementation(fn: T, callback: () => void): void;
   mockName(name: string): this;
   mockReturnThis(): this;
   mockReturnValue(value: ReturnType<T>): this;
@@ -768,6 +770,38 @@ export class ModuleMocker {
         return f;
       };
 
+      f.withImplementation = withImplementation.bind(this);
+
+      function withImplementation(fn: T, callback: () => void): void;
+      function withImplementation(
+        fn: T,
+        callback: () => Promise<unknown>,
+      ): Promise<void>;
+      function withImplementation(
+        this: ModuleMocker,
+        fn: T,
+        callback: (() => void) | (() => Promise<unknown>),
+      ): void | Promise<void> {
+        // Remember previous mock implementation, then set new one
+        const mockConfig = this._ensureMockConfig(f);
+        const previousImplementation = mockConfig.mockImpl;
+        mockConfig.mockImpl = fn;
+
+        const returnedValue = callback();
+
+        if (
+          returnedValue != null &&
+          typeof returnedValue === 'object' &&
+          typeof returnedValue.then === 'function'
+        ) {
+          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);