From 30508851d7ab336f1ff265f7a19e2450df8ffcaa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jesper=20S=C3=B8rensen?= Date: Fri, 6 Dec 2019 11:51:48 +0000 Subject: [PATCH] feat: implement `mockRestoreInitialImplementation` and `restoreAllInitialMockImplementations` --- packages/jest-environment/src/index.ts | 9 ++++ .../jest-mock/src/__tests__/index.test.ts | 43 +++++++++++++++++++ packages/jest-mock/src/index.ts | 26 +++++++++-- packages/jest-runtime/src/index.ts | 9 ++++ 4 files changed, 84 insertions(+), 3 deletions(-) diff --git a/packages/jest-environment/src/index.ts b/packages/jest-environment/src/index.ts index 5f265aba0b65..227971a0110d 100644 --- a/packages/jest-environment/src/index.ts +++ b/packages/jest-environment/src/index.ts @@ -190,6 +190,15 @@ export interface Jest { * useful to isolate modules where local state might conflict between tests. */ resetModules(): Jest; + /** + * Restores all mocks back to the initial mock implementation (or none if no + * implementation has been set) and clears the mock. Equivalent to calling + * `.mockRestoreInitialImplementation` on every mocked function. + * + * Beware that jest.restoreAllInitialMockImplementations() will only restore the + * initial mocked implementation, it will not restore the original implementation. + */ + restoreAllInitialMockImplementations(): Jest; /** * Restores all mocks back to their original value. Equivalent to calling * `.mockRestore` on every mocked function. diff --git a/packages/jest-mock/src/__tests__/index.test.ts b/packages/jest-mock/src/__tests__/index.test.ts index df330d6354e1..ef0b824a7dda 100644 --- a/packages/jest-mock/src/__tests__/index.test.ts +++ b/packages/jest-mock/src/__tests__/index.test.ts @@ -536,6 +536,49 @@ describe('moduleMocker', () => { expect(fn2()).not.toEqual('abcd'); }); + it('supports restoring the initial implementation', () => { + const fn = moduleMocker.fn(); + fn.mockImplementation(() => 'initial result'); + + const initial = fn(); + expect(initial).toEqual('initial result'); + + fn.mockImplementation(() => 'overridden result'); + const overridden = fn(); + expect(overridden).toEqual('overridden result'); + + fn.mockRestoreInitialImplementation(); + const restored = fn(); + expect(restored).toEqual('initial result'); + }); + + it('supports restoring all initial mock implementations', () => { + const fn1 = moduleMocker.fn(); + fn1.mockImplementation(() => 'fn1-before'); + expect(fn1(1, 2, 3)).toEqual('fn1-before'); + expect(fn1).toHaveBeenCalledWith(1, 2, 3); + + fn1.mockImplementation(() => 'fn1-after'); + expect(fn1('a', 'b', 'c')).toEqual('fn1-after'); + expect(fn1).toHaveBeenCalledWith('a', 'b', 'c'); + + const fn2 = moduleMocker.fn(); + fn2.mockImplementation(() => 'fn2-before'); + expect(fn2(3, 2, 1)).toEqual('fn2-before'); + expect(fn2).toHaveBeenCalledWith(3, 2, 1); + + fn2.mockImplementation(() => 'fn2-after'); + expect(fn2('a', 'b', 'c')).toEqual('fn2-after'); + expect(fn2).toHaveBeenCalledWith('a', 'b', 'c'); + + moduleMocker.restoreAllInitialMockImplementations(); + expect(fn1).not.toHaveBeenCalled(); + expect(fn2).not.toHaveBeenCalled(); + + expect(fn1()).toEqual('fn1-before'); + expect(fn2()).toEqual('fn2-before'); + }); + it('maintains function arity', () => { const mockFunctionArity1 = moduleMocker.fn(x => x); const mockFunctionArity2 = moduleMocker.fn((x, y) => y); diff --git a/packages/jest-mock/src/index.ts b/packages/jest-mock/src/index.ts index 8afcd60b6a56..6b8f5a306567 100644 --- a/packages/jest-mock/src/index.ts +++ b/packages/jest-mock/src/index.ts @@ -71,6 +71,7 @@ type MockFunctionState> = { }; type MockFunctionConfig = { + initialImpl: Function | undefined; mockImpl: Function | undefined; mockName: string; specificReturnValues: Array; @@ -105,6 +106,7 @@ interface MockInstance> { mockClear(): this; mockReset(): this; mockRestore(): void; + mockRestoreInitialImplementation: () => void; mockImplementation(fn: (...args: Y) => T): this; mockImplementation(fn: () => Promise): this; mockImplementationOnce(fn: (...args: Y) => T): this; @@ -364,7 +366,7 @@ function isReadonlyProp(object: any, prop: string): boolean { class ModuleMockerClass { private _environmentGlobal: Global; private _mockState: WeakMap, MockFunctionState>; - private _mockConfigRegistry: WeakMap; + private _mockConfigRegistry: Map, MockFunctionConfig>; private _spyState: Set<() => void>; private _invocationCallCounter: number; ModuleMocker: typeof ModuleMockerClass; @@ -377,7 +379,7 @@ class ModuleMockerClass { constructor(global: Global) { this._environmentGlobal = global; this._mockState = new WeakMap(); - this._mockConfigRegistry = new WeakMap(); + this._mockConfigRegistry = new Map(); this._spyState = new Set(); this.ModuleMocker = ModuleMockerClass; this._invocationCallCounter = 1; @@ -454,6 +456,7 @@ class ModuleMockerClass { private _defaultMockConfig(): MockFunctionConfig { return { + initialImpl: undefined, mockImpl: undefined, mockName: 'jest.fn()', specificMockImpls: [], @@ -652,6 +655,16 @@ class ModuleMockerClass { return restore ? restore() : undefined; }; + f.mockRestoreInitialImplementation = () => { + f.mockClear(); + + const currentConfig = this._ensureMockConfig(f); + const newConfig = this._defaultMockConfig(); + newConfig.initialImpl = newConfig.mockImpl = currentConfig.initialImpl; + + this._mockConfigRegistry.set(f, newConfig); + }; + f.mockReturnValueOnce = (value: T) => // next function call will return this value or default return value f.mockImplementationOnce(() => value); @@ -687,6 +700,7 @@ class ModuleMockerClass { ): Mock => { // next function call will use mock implementation return value const mockConfig = this._ensureMockConfig(f); + if (!mockConfig.initialImpl) mockConfig.initialImpl = fn; mockConfig.mockImpl = fn; return f; }; @@ -1076,7 +1090,7 @@ class ModuleMockerClass { } resetAllMocks() { - this._mockConfigRegistry = new WeakMap(); + this._mockConfigRegistry = new Map(); this._mockState = new WeakMap(); } @@ -1085,6 +1099,12 @@ class ModuleMockerClass { this._spyState = new Set(); } + restoreAllInitialMockImplementations() { + for (const f of this._mockConfigRegistry.keys()) { + f.mockRestoreInitialImplementation(); + } + } + private _typeOf(value: any): string { return value == null ? '' + value : typeof value; } diff --git a/packages/jest-runtime/src/index.ts b/packages/jest-runtime/src/index.ts index a082739ba7fb..30e51a88a19b 100644 --- a/packages/jest-runtime/src/index.ts +++ b/packages/jest-runtime/src/index.ts @@ -614,6 +614,10 @@ class Runtime { this._moduleMocker.clearAllMocks(); } + restoreAllInitialMockImplementations() { + this._moduleMocker.restoreAllInitialMockImplementations(); + } + private _resolveModule(from: Config.Path, to?: string) { return to ? this._resolver.resolveModule(from, to) : from; } @@ -1018,6 +1022,10 @@ class Runtime { this.restoreAllMocks(); return jestObject; }; + const restoreAllInitialMockImplementations = () => { + this.restoreAllInitialMockImplementations(); + return jestObject; + }; const useFakeTimers = () => { _getFakeTimers().useFakeTimers(); return jestObject; @@ -1093,6 +1101,7 @@ class Runtime { resetAllMocks, resetModuleRegistry: resetModules, resetModules, + restoreAllInitialMockImplementations, restoreAllMocks, retryTimes, runAllImmediates: () => _getFakeTimers().runAllImmediates(),