Skip to content

Commit

Permalink
feat: add onGenerateMock transformer callback (#15429) (#15433)
Browse files Browse the repository at this point in the history
  • Loading branch information
MillerSvt authored Jan 15, 2025
1 parent 0630928 commit 3088367
Show file tree
Hide file tree
Showing 5 changed files with 152 additions and 1 deletion.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
45 changes: 45 additions & 0 deletions docs/JestObjectAPI.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
10 changes: 10 additions & 0 deletions packages/jest-environment/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<T>(cb: (moduleName: string, moduleMock: T) => T): Jest;
/**
* Replaces property on an object with another value.
*
Expand Down
78 changes: 78 additions & 0 deletions packages/jest-runtime/src/__tests__/runtime_mock.test.js
Original file line number Diff line number Diff line change
Expand Up @@ -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,
});
});
});
});
19 changes: 18 additions & 1 deletion packages/jest-runtime/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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),
Expand Down

0 comments on commit 3088367

Please sign in to comment.