feat: add onGenerateMock transformer callback (#15429)

Author: MillerSvt

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
 

docs/JestObjectAPI.md

@@ -528,6 +528,51 @@ getRandom(); // Always returns 10
 
 Returns a mock module instead of the actual module, bypassing all checks on whether the module should be required normally or not.
 
+### `jest.onGenerateMock(cb)`
+
+Registers a callback function that is invoked whenever Jest generates a mock for a module.
+This callback allows you to modify the mock before it is returned to the rest of your tests.
+
+Parameters for callback:
+1. `moduleName: string` - The name of the module that is being mocked.
+2. `moduleMock: T` - The mock object that Jest has generated for the module. This object can be modified or replaced before returning.
+
+Behaviour:
+
+- If multiple callbacks are registered via consecutive `onGenerateMock` calls, they will be invoked **in the order they were added**.
+- Each callback receives the output of the previous callback as its `moduleMock`. This makes it possible to apply multiple layers of transformations to the same mock.
+
+```js
+jest.onGenerateMock((moduleName, moduleMock) => {
+  // Inspect the module name and decide how to transform the mock
+  if (moduleName.includes('Database')) {
+    // For demonstration, let's replace a method with our own custom mock
+    moduleMock.connect = jest.fn().mockImplementation(() => {
+      console.log('Connected to mock DB');
+    });
+  }
+
+  // Return the (potentially modified) mock
+  return moduleMock;
+});
+
+// Apply mock for module
+jest.mock('./Database');
+
+// Later in your tests
+import Database from './Database';
+// The `Database` mock now has any transformations applied by our callback
+```
+
+:::note
+
+The `onGenerateMock` callback is not called for manually created mocks, such as:
+
+- Mocks defined in a `__mocks__` folder
+- Explicit factories provided via `jest.mock('moduleName', () => { ... })`
+
+:::
+
 ### `jest.resetModules()`
 
 Resets the module registry - the cache of all required modules. This is useful to isolate modules where local state might conflict between tests.

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),