feat: add onGenerateMock transformer callback (jestjs#15429)

Author: s.v.zaytsev

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__/__snapshots__/runtime_mock.test.js.snap

@@ -0,0 +1,80 @@
+// Jest Snapshot v1, https://goo.gl/fbAQLP
+
+exports[`Runtime jest.onGenerateMock calls single callback and returns transformed value 1`] = `
+Object {
+  "filename": "/Users/s.v.zaytsev/IdeaProjects/jest/packages/jest-runtime/src/__tests__/test_root/RegularModule.js",
+  "getModuleStateValue": [MockFunction],
+  "isLoaded": [MockFunction],
+  "isRealModule": true,
+  "jest": Object {
+    "advanceTimersByTime": [MockFunction],
+    "advanceTimersByTimeAsync": [MockFunction],
+    "advanceTimersToNextFrame": [MockFunction],
+    "advanceTimersToNextTimer": [MockFunction],
+    "advanceTimersToNextTimerAsync": [MockFunction],
+    "autoMockOff": [MockFunction],
+    "autoMockOn": [MockFunction],
+    "clearAllMocks": [MockFunction],
+    "clearAllTimers": [MockFunction],
+    "createMockFromModule": [MockFunction],
+    "deepUnmock": [MockFunction],
+    "disableAutomock": [MockFunction],
+    "doMock": [MockFunction],
+    "dontMock": [MockFunction],
+    "enableAutomock": [MockFunction],
+    "fn": [MockFunction],
+    "getRealSystemTime": [MockFunction],
+    "getSeed": [MockFunction],
+    "getTimerCount": [MockFunction],
+    "isEnvironmentTornDown": [MockFunction],
+    "isMockFunction": [MockFunction],
+    "isolateModules": [MockFunction],
+    "isolateModulesAsync": [MockFunction],
+    "mock": [MockFunction],
+    "mocked": [MockFunction],
+    "now": [MockFunction],
+    "onGenerateMock": [MockFunction],
+    "replaceProperty": [MockFunction],
+    "requireActual": [MockFunction],
+    "requireMock": [MockFunction],
+    "resetAllMocks": [MockFunction],
+    "resetModules": [MockFunction],
+    "restoreAllMocks": [MockFunction],
+    "retryTimes": [MockFunction],
+    "runAllImmediates": [MockFunction],
+    "runAllTicks": [MockFunction],
+    "runAllTimers": [MockFunction],
+    "runAllTimersAsync": [MockFunction],
+    "runOnlyPendingTimers": [MockFunction],
+    "runOnlyPendingTimersAsync": [MockFunction],
+    "setMock": [MockFunction],
+    "setSystemTime": [MockFunction],
+    "setTimeout": [MockFunction],
+    "spyOn": [MockFunction],
+    "unmock": [MockFunction],
+    "unstable_mockModule": [MockFunction],
+    "unstable_unmockModule": [MockFunction],
+    "useFakeTimers": [MockFunction],
+    "useRealTimers": [MockFunction],
+  },
+  "lazyRequire": [MockFunction],
+  "loaded": false,
+  "module": Object {
+    "children": Array [],
+    "exports": [Circular],
+    "filename": "/Users/s.v.zaytsev/IdeaProjects/jest/packages/jest-runtime/src/__tests__/test_root/RegularModule.js",
+    "id": "/Users/s.v.zaytsev/IdeaProjects/jest/packages/jest-runtime/src/__tests__/test_root/RegularModule.js",
+    "isPreloading": false,
+    "loaded": true,
+    "main": null,
+    "path": "/Users/s.v.zaytsev/IdeaProjects/jest/packages/jest-runtime/src/__tests__/test_root",
+    "paths": Array [],
+    "require": [MockFunction],
+  },
+  "object": Object {},
+  "parent": null,
+  "path": "/Users/s.v.zaytsev/IdeaProjects/jest/packages/jest-runtime/src/__tests__/test_root",
+  "paths": Array [],
+  "setModuleStateValue": [MockFunction],
+}
+`;

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

@@ -136,4 +136,83 @@ 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(),
+      );
+      expect(onGenerateMock.mock.calls[0][1]).toMatchSnapshot();
+
+      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),