feat: add onGenerateMock transformer callback (#15429)

CHANGELOG.md

@@ -46,6 +46,7 @@
 - `[@jest/util-snapshot]` Extract utils used by tooling from `jest-snapshot` into its own package ([#15095](https://github.com/facebook/jest/pull/15095))
 - `[pretty-format]` [**BREAKING**] Do not render empty string children (`''`) in React plugin ([#14470](https://github.com/facebook/jest/pull/14470))
 - `[jest-each]` Introduce `%$` option to add number of the test to its title ([#14710](https://github.com/jestjs/jest/pull/14710))
+- `[jest-runtime]` Add `onGenerateMock` transformer callback for auto generated callbacks ([#15433](https://github.com/jestjs/jest/pull/15433))
 
 ### Fixes
 

packages/jest-environment/src/index.ts

@@ -207,6 +207,18 @@ export interface Jest {
     moduleFactory?: () => T,
     options?: {virtual?: boolean},
   ): Jest;
+  /**
+   * Registers a callback function that is invoked whenever a mock is generated for a module.
+   * This callback is passed the module name and the newly created mock object, and must return
+   * the (potentially modified) mock object.
+   *
+   * If multiple callbacks are registered, they will be called in the order they were added.
+   * Each callback receives the result of the previous callback as the `moduleMock` parameter,
+   * making it possible to apply sequential transformations.
+   *
+   * @param cb
+   */
+  onGenerateMock<T>(cb: (moduleName: string, moduleMock: T) => T): Jest;
   /**
    * Mocks a module with the provided module factory when it is being imported.
    */

packages/jest-runtime/src/__tests__/runtime_mock.test.js

@@ -136,4 +136,82 @@ describe('Runtime', () => {
       ).toBe(mockReference);
     });
   });
+
+  describe('jest.onGenerateMock', () => {
+    it('calls single callback and returns transformed value', async () => {
+      const runtime = await createRuntime(__filename);
+      const mockReference = {isMock: true};
+      const root = runtime.requireModule(runtime.__mockRootPath, rootJsPath);
+      // Erase module registry because root.js requires most other modules.
+      root.jest.resetModules();
+
+      const onGenerateMock = jest.fn((moduleName, moduleMock) => mockReference);
+
+      root.jest.onGenerateMock(onGenerateMock);
+      root.jest.mock('RegularModule');
+      root.jest.mock('ManuallyMocked');
+
+      expect(
+        runtime.requireModuleOrMock(runtime.__mockRootPath, 'RegularModule'),
+      ).toEqual(mockReference);
+      expect(onGenerateMock).toHaveBeenCalledWith(
+        'RegularModule',
+        expect.anything(),
+      );
+
+      onGenerateMock.mockReset();
+
+      expect(
+        runtime.requireModuleOrMock(runtime.__mockRootPath, 'ManuallyMocked'),
+      ).not.toEqual(mockReference);
+      expect(onGenerateMock).not.toHaveBeenCalled();
+    });
+
+    it('calls multiple callback and returns transformed value', async () => {
+      const runtime = await createRuntime(__filename);
+      const root = runtime.requireModule(runtime.__mockRootPath, rootJsPath);
+      // Erase module registry because root.js requires most other modules.
+      root.jest.resetModules();
+
+      const onGenerateMock1 = jest.fn((moduleName, moduleMock) => ({
+        isMock: true,
+        value: 1,
+      }));
+
+      const onGenerateMock2 = jest.fn((moduleName, moduleMock) => ({
+        ...moduleMock,
+        value: moduleMock.value + 1,
+      }));
+
+      const onGenerateMock3 = jest.fn((moduleName, moduleMock) => ({
+        ...moduleMock,
+        value: moduleMock.value ** 2,
+      }));
+
+      root.jest.onGenerateMock(onGenerateMock1);
+      root.jest.onGenerateMock(onGenerateMock2);
+      root.jest.onGenerateMock(onGenerateMock3);
+      root.jest.mock('RegularModule');
+      root.jest.mock('ManuallyMocked');
+
+      expect(
+        runtime.requireModuleOrMock(runtime.__mockRootPath, 'RegularModule'),
+      ).toEqual({
+        isMock: true,
+        value: 4,
+      });
+      expect(onGenerateMock1).toHaveBeenCalledWith(
+        'RegularModule',
+        expect.anything(),
+      );
+      expect(onGenerateMock2).toHaveBeenCalledWith('RegularModule', {
+        isMock: true,
+        value: 1,
+      });
+      expect(onGenerateMock3).toHaveBeenCalledWith('RegularModule', {
+        isMock: true,
+        value: 2,
+      });
+    });
+  });
 });

packages/jest-runtime/src/index.ts

@@ -171,6 +171,9 @@ export default class Runtime {
   private readonly _environment: JestEnvironment;
   private readonly _explicitShouldMock: Map<string, boolean>;
   private readonly _explicitShouldMockModule: Map<string, boolean>;
+  private readonly _onGenerateMock: Set<
+    (moduleName: string, moduleMock: any) => any
+  >;
   private _fakeTimersImplementation:
     | LegacyFakeTimers<unknown>
     | ModernFakeTimers
@@ -235,6 +238,7 @@ export default class Runtime {
     this._globalConfig = globalConfig;
     this._explicitShouldMock = new Map();
     this._explicitShouldMockModule = new Map();
+    this._onGenerateMock = new Set();
     this._internalModuleRegistry = new Map();
     this._isCurrentlyExecutingManualMock = null;
     this._mainModule = null;
@@ -1930,10 +1934,16 @@ export default class Runtime {
       }
       this._mockMetaDataCache.set(modulePath, mockMetadata);
     }
-    return this._moduleMocker.generateFromMetadata<T>(
+    let moduleMock = this._moduleMocker.generateFromMetadata<T>(
       // added above if missing
       this._mockMetaDataCache.get(modulePath)!,
     );
+
+    for (const onGenerateMock of this._onGenerateMock) {
+      moduleMock = onGenerateMock(moduleName, moduleMock);
+    }
+
+    return moduleMock;
   }
 
   private _shouldMockCjs(
@@ -2193,6 +2203,12 @@ export default class Runtime {
       this._explicitShouldMock.set(moduleID, true);
       return jestObject;
     };
+    const onGenerateMock: Jest['onGenerateMock'] = <T>(
+      cb: (moduleName: string, moduleMock: T) => T,
+    ) => {
+      this._onGenerateMock.add(cb);
+      return jestObject;
+    };
     const setMockFactory = (
       moduleName: string,
       mockFactory: () => unknown,
@@ -2364,6 +2380,7 @@ export default class Runtime {
       mock,
       mocked,
       now: () => _getFakeTimers().now(),
+      onGenerateMock,
       replaceProperty,
       requireActual: moduleName => this.requireActual(from, moduleName),
       requireMock: moduleName => this.requireMock(from, moduleName),