fix(jest-mock): restore withImplementation after errors

CHANGELOG.md

@@ -11,6 +11,7 @@
 - `[expect, jest-message-util, jest-pattern, jest-regex-util, jest-util]` Revert `node:` protocol imports to restore webpack/browser-bundle compatibility ([#16167](https://github.com/jestjs/jest/pull/16167))
 - `[expect]` Widen `toMatchObject` and `objectContaining` parameter type from `Record<string, unknown>` to `object` so class instances are accepted ([#16196](https://github.com/jestjs/jest/pull/16196))
 - `[@jest-environment/jsdom-abstract]` Make `@types/jsdom` a peer dependency ([#16166](https://github.com/jestjs/jest/pull/16166))
+- `[jest-mock]` Restore `withImplementation(...)` after sync throws and async rejections ([#16219](https://github.com/jestjs/jest/pull/16219))
 - `[jest-runtime]` Fall back to native ESM when a `.js` file contains ESM syntax but has no `"type":"module"` marker ([#16152](https://github.com/jestjs/jest/pull/16152))
 - `[jest-runtime]` Support older test environments whose `moduleMocker` does not implement `clearMocksOnScope` ([#16169](https://github.com/jestjs/jest/pull/16169))
 

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

@@ -1166,6 +1166,43 @@ describe('moduleMocker', () => {
       expect.assertions(4);
     });
 
+    it('restores the previous implementation after the callback throws', () => {
+      const mock = jest
+        .fn(() => 'outside callback')
+        .mockImplementationOnce(() => 'once');
+
+      expect(() =>
+        mock.withImplementation(
+          () => 'inside callback',
+          () => {
+            expect(mock()).toBe('inside callback');
+            throw new Error('boom');
+          },
+        ),
+      ).toThrow('boom');
+
+      expect(mock()).toBe('once');
+      expect(mock()).toBe('outside callback');
+    });
+
+    it('restores whenCalledWith fallbacks after the callback rejects', async () => {
+      const mock = jest.fn<(arg?: string) => string>(() => 'outside callback');
+      mock.whenCalledWith('branch').mockReturnValue('branch');
+
+      await expect(
+        mock.withImplementation(
+          () => 'inside callback',
+          async () => {
+            expect(mock('other')).toBe('inside callback');
+            throw new Error('boom');
+          },
+        ),
+      ).rejects.toThrow('boom');
+
+      expect(mock('branch')).toBe('branch');
+      expect(mock('other')).toBe('outside callback');
+    });
+
     it('mockImplementationOnce does not bleed into withImplementation', () => {
       const mock = jest
         .fn(() => 'outside callback')

packages/jest-mock/src/index.ts

@@ -910,22 +910,35 @@ export class ModuleMocker {
         const previousImplementation = mockConfig.mockImpl;
         const previousSpecificImplementations = mockConfig.specificMockImpls;
         const previousFallbackImpl = mockConfig.fallbackImpl;
+        const restore = () => {
+          mockConfig.mockImpl = previousImplementation;
+          mockConfig.specificMockImpls = previousSpecificImplementations;
+          mockConfig.fallbackImpl = previousFallbackImpl;
+        };
         mockConfig.mockImpl = fn;
         mockConfig.specificMockImpls = [];
 
-        const returnedValue = callback();
+        let returnedValue: ReturnType<typeof callback>;
+        try {
+          returnedValue = callback();
+        } catch (error) {
+          restore();
+          throw error;
+        }
 
         if (isPromise(returnedValue)) {
-          return returnedValue.then(() => {
-            mockConfig.mockImpl = previousImplementation;
-            mockConfig.specificMockImpls = previousSpecificImplementations;
-            mockConfig.fallbackImpl = previousFallbackImpl;
-          });
-        } else {
-          mockConfig.mockImpl = previousImplementation;
-          mockConfig.specificMockImpls = previousSpecificImplementations;
-          mockConfig.fallbackImpl = previousFallbackImpl;
+          return this._environmentGlobal.Promise.resolve(returnedValue).then(
+            () => {
+              restore();
+            },
+            error => {
+              restore();
+              throw error;
+            },
+          );
         }
+
+        restore();
       }
 
       f.mockImplementation = (fn: T) => {