feat: implement mockRestoreInitialImplementation and restoreAllInitialMockImplementations

packages/jest-environment/src/index.ts

@@ -190,6 +190,15 @@ export interface Jest {
    * useful to isolate modules where local state might conflict between tests.
    */
   resetModules(): Jest;
+  /**
+   * Restores all mocks back to the initial mock implementation (or none if no
+   * implementation has been set) and clears the mock. Equivalent to calling
+   * `.mockRestoreInitialImplementation` on every mocked function.
+   *
+   * Beware that jest.restoreAllInitialMockImplementations() will only restore the
+   * initial mocked implementation, it will not restore the original implementation.
+   */
+  restoreAllInitialMockImplementations(): Jest;
   /**
    * Restores all mocks back to their original value. Equivalent to calling
    * `.mockRestore` on every mocked function.

packages/jest-mock/src/__tests__/index.test.ts

@@ -536,6 +536,49 @@ describe('moduleMocker', () => {
         expect(fn2()).not.toEqual('abcd');
       });
 
+      it('supports restoring the initial implementation', () => {
+        const fn = moduleMocker.fn();
+        fn.mockImplementation(() => 'initial result');
+
+        const initial = fn();
+        expect(initial).toEqual('initial result');
+
+        fn.mockImplementation(() => 'overridden result');
+        const overridden = fn();
+        expect(overridden).toEqual('overridden result');
+
+        fn.mockRestoreInitialImplementation();
+        const restored = fn();
+        expect(restored).toEqual('initial result');
+      });
+
+      it('supports restoring all initial mock implementations', () => {
+        const fn1 = moduleMocker.fn();
+        fn1.mockImplementation(() => 'fn1-before');
+        expect(fn1(1, 2, 3)).toEqual('fn1-before');
+        expect(fn1).toHaveBeenCalledWith(1, 2, 3);
+
+        fn1.mockImplementation(() => 'fn1-after');
+        expect(fn1('a', 'b', 'c')).toEqual('fn1-after');
+        expect(fn1).toHaveBeenCalledWith('a', 'b', 'c');
+
+        const fn2 = moduleMocker.fn();
+        fn2.mockImplementation(() => 'fn2-before');
+        expect(fn2(3, 2, 1)).toEqual('fn2-before');
+        expect(fn2).toHaveBeenCalledWith(3, 2, 1);
+
+        fn2.mockImplementation(() => 'fn2-after');
+        expect(fn2('a', 'b', 'c')).toEqual('fn2-after');
+        expect(fn2).toHaveBeenCalledWith('a', 'b', 'c');
+
+        moduleMocker.restoreAllInitialMockImplementations();
+        expect(fn1).not.toHaveBeenCalled();
+        expect(fn2).not.toHaveBeenCalled();
+
+        expect(fn1()).toEqual('fn1-before');
+        expect(fn2()).toEqual('fn2-before');
+      });
+
       it('maintains function arity', () => {
         const mockFunctionArity1 = moduleMocker.fn(x => x);
         const mockFunctionArity2 = moduleMocker.fn((x, y) => y);

packages/jest-mock/src/index.ts

@@ -71,6 +71,7 @@ type MockFunctionState<T, Y extends Array<unknown>> = {
 };
 
 type MockFunctionConfig = {
+  initialImpl: Function | undefined;
   mockImpl: Function | undefined;
   mockName: string;
   specificReturnValues: Array<unknown>;
@@ -105,6 +106,7 @@ interface MockInstance<T, Y extends Array<unknown>> {
   mockClear(): this;
   mockReset(): this;
   mockRestore(): void;
+  mockRestoreInitialImplementation: () => void;
   mockImplementation(fn: (...args: Y) => T): this;
   mockImplementation(fn: () => Promise<T>): this;
   mockImplementationOnce(fn: (...args: Y) => T): this;
@@ -364,7 +366,7 @@ function isReadonlyProp(object: any, prop: string): boolean {
 class ModuleMockerClass {
   private _environmentGlobal: Global;
   private _mockState: WeakMap<Mock<any, any>, MockFunctionState<any, any>>;
-  private _mockConfigRegistry: WeakMap<Function, MockFunctionConfig>;
+  private _mockConfigRegistry: Map<Mock<any, any>, MockFunctionConfig>;
   private _spyState: Set<() => void>;
   private _invocationCallCounter: number;
   ModuleMocker: typeof ModuleMockerClass;
@@ -377,7 +379,7 @@ class ModuleMockerClass {
   constructor(global: Global) {
     this._environmentGlobal = global;
     this._mockState = new WeakMap();
-    this._mockConfigRegistry = new WeakMap();
+    this._mockConfigRegistry = new Map();
     this._spyState = new Set();
     this.ModuleMocker = ModuleMockerClass;
     this._invocationCallCounter = 1;
@@ -454,6 +456,7 @@ class ModuleMockerClass {
 
   private _defaultMockConfig(): MockFunctionConfig {
     return {
+      initialImpl: undefined,
       mockImpl: undefined,
       mockName: 'jest.fn()',
       specificMockImpls: [],
@@ -652,6 +655,16 @@ class ModuleMockerClass {
         return restore ? restore() : undefined;
       };
 
+      f.mockRestoreInitialImplementation = () => {
+        f.mockClear();
+
+        const currentConfig = this._ensureMockConfig(f);
+        const newConfig = this._defaultMockConfig();
+        newConfig.initialImpl = newConfig.mockImpl = currentConfig.initialImpl;
+
+        this._mockConfigRegistry.set(f, newConfig);
+      };
+
       f.mockReturnValueOnce = (value: T) =>
         // next function call will return this value or default return value
         f.mockImplementationOnce(() => value);
@@ -687,6 +700,7 @@ class ModuleMockerClass {
       ): Mock<T, Y> => {
         // next function call will use mock implementation return value
         const mockConfig = this._ensureMockConfig(f);
+        if (!mockConfig.initialImpl) mockConfig.initialImpl = fn;
         mockConfig.mockImpl = fn;
         return f;
       };
@@ -1076,7 +1090,7 @@ class ModuleMockerClass {
   }
 
   resetAllMocks() {
-    this._mockConfigRegistry = new WeakMap();
+    this._mockConfigRegistry = new Map();
     this._mockState = new WeakMap();
   }
 
@@ -1085,6 +1099,12 @@ class ModuleMockerClass {
     this._spyState = new Set();
   }
 
+  restoreAllInitialMockImplementations() {
+    for (const f of this._mockConfigRegistry.keys()) {
+      f.mockRestoreInitialImplementation();
+    }
+  }
+
   private _typeOf(value: any): string {
     return value == null ? '' + value : typeof value;
   }

packages/jest-runtime/src/index.ts

@@ -614,6 +614,10 @@ class Runtime {
     this._moduleMocker.clearAllMocks();
   }
 
+  restoreAllInitialMockImplementations() {
+    this._moduleMocker.restoreAllInitialMockImplementations();
+  }
+
   private _resolveModule(from: Config.Path, to?: string) {
     return to ? this._resolver.resolveModule(from, to) : from;
   }
@@ -1018,6 +1022,10 @@ class Runtime {
       this.restoreAllMocks();
       return jestObject;
     };
+    const restoreAllInitialMockImplementations = () => {
+      this.restoreAllInitialMockImplementations();
+      return jestObject;
+    };
     const useFakeTimers = () => {
       _getFakeTimers().useFakeTimers();
       return jestObject;
@@ -1093,6 +1101,7 @@ class Runtime {
       resetAllMocks,
       resetModuleRegistry: resetModules,
       resetModules,
+      restoreAllInitialMockImplementations,
       restoreAllMocks,
       retryTimes,
       runAllImmediates: () => _getFakeTimers().runAllImmediates(),