From 3088367874ada4f6bfe9f08fc0cecbf969de2eb5 Mon Sep 17 00:00:00 2001 From: Svyatoslav Zaytsev <143048525+MillerSvt@users.noreply.github.com> Date: Wed, 15 Jan 2025 15:24:14 +0500 Subject: [PATCH] feat: add onGenerateMock transformer callback (#15429) (#15433) --- CHANGELOG.md | 1 + docs/JestObjectAPI.md | 45 +++++++++++ packages/jest-environment/src/index.ts | 10 +++ .../src/__tests__/runtime_mock.test.js | 78 +++++++++++++++++++ packages/jest-runtime/src/index.ts | 19 ++++- 5 files changed, 152 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 695566a72df8..10578044269e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -47,6 +47,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 diff --git a/docs/JestObjectAPI.md b/docs/JestObjectAPI.md index 8b0e6ece428f..66a292876da2 100644 --- a/docs/JestObjectAPI.md +++ b/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. diff --git a/packages/jest-environment/src/index.ts b/packages/jest-environment/src/index.ts index 9f1eccc3c378..04a02ae10c0c 100644 --- a/packages/jest-environment/src/index.ts +++ b/packages/jest-environment/src/index.ts @@ -225,6 +225,16 @@ export interface Jest { * Returns the current time in ms of the fake timer clock. */ now(): number; + /** + * 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. + */ + onGenerateMock(cb: (moduleName: string, moduleMock: T) => T): Jest; /** * Replaces property on an object with another value. * diff --git a/packages/jest-runtime/src/__tests__/runtime_mock.test.js b/packages/jest-runtime/src/__tests__/runtime_mock.test.js index bec7d128c2b3..628af52ca21a 100644 --- a/packages/jest-runtime/src/__tests__/runtime_mock.test.js +++ b/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, + }); + }); + }); }); diff --git a/packages/jest-runtime/src/index.ts b/packages/jest-runtime/src/index.ts index 477be4b7ef22..e054e894d211 100644 --- a/packages/jest-runtime/src/index.ts +++ b/packages/jest-runtime/src/index.ts @@ -171,6 +171,9 @@ export default class Runtime { private readonly _environment: JestEnvironment; private readonly _explicitShouldMock: Map; private readonly _explicitShouldMockModule: Map; + private readonly _onGenerateMock: Set< + (moduleName: string, moduleMock: any) => any + >; private _fakeTimersImplementation: | LegacyFakeTimers | 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( + let moduleMock = this._moduleMocker.generateFromMetadata( // 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'] = ( + 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),