diff --git a/docs/api/vi.md b/docs/api/vi.md index 1a4772dfd..64743a0bc 100644 --- a/docs/api/vi.md +++ b/docs/api/vi.md @@ -459,6 +459,19 @@ expect(spy).toHaveBeenCalled() expect(spy).toHaveReturnedWith(1) ``` +::: tip +In environments that support [Explicit Resource Management](https://github.com/tc39/proposal-explicit-resource-management), you can use `using` instead of `const` to automatically call `mockRestore` on any mocked function when the containing block is exited. This is especially useful for spied methods: + +```ts +it('calls console.log', () => { + using spy = vi.spyOn(console, 'log').mockImplementation(() => {}) + debug('message') + expect(spy).toHaveBeenCalled() +}) +// console.log is restored here +``` +::: + ::: tip You can call [`vi.restoreAllMocks`](#vi-restoreallmocks) inside [`afterEach`](/api/#aftereach) (or enable [`test.restoreMocks`](/config/#restoreMocks)) to restore all methods to their original implementations. This will restore the original [object descriptor](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Object/defineProperty), so you won't be able to change method's implementation: diff --git a/packages/spy/src/index.ts b/packages/spy/src/index.ts index f6b0e8386..1a2360da7 100644 --- a/packages/spy/src/index.ts +++ b/packages/spy/src/index.ts @@ -175,7 +175,7 @@ Jest uses the latter for `MockInstance.mockImplementation` etc... and it allows const boolFn: Jest.Mock<() => boolean> = jest.fn<() => true>(() => true) */ /* eslint-disable ts/method-signature-style */ -export interface MockInstance { +export interface MockInstance extends Disposable { /** * Use it to return the name assigned to the mock with the `.mockName(name)` method. By default, it will return `vi.fn()`. * @see https://vitest.dev/api/mock#getmockname @@ -549,6 +549,10 @@ function enhanceSpy( return stub } + if (Symbol.dispose) { + stub[Symbol.dispose] = () => stub.mockRestore() + } + stub.getMockImplementation = () => implementationChangedTemporarily ? implementation : (onceImplementations.at(0) || implementation) stub.mockImplementation = (fn: T) => { diff --git a/test/core/test/jest-mock.test.ts b/test/core/test/jest-mock.test.ts index 9276335e4..a4a01f7b5 100644 --- a/test/core/test/jest-mock.test.ts +++ b/test/core/test/jest-mock.test.ts @@ -560,6 +560,30 @@ describe('jest mock compat layer', () => { expect(fn.getMockImplementation()).toBe(temporaryMockImplementation) }) + describe('is disposable', () => { + describe.runIf(Symbol.dispose)('in environments supporting it', () => { + it('has dispose property', () => { + expect(vi.fn()[Symbol.dispose]).toBeTypeOf('function') + }) + it('calls mockRestore when disposing', () => { + const fn = vi.fn() + const restoreSpy = vi.spyOn(fn, 'mockRestore') + { + using _fn2 = fn + } + expect(restoreSpy).toHaveBeenCalled() + }) + it('allows disposal when using mockImplementation', () => { + expect(vi.fn().mockImplementation(() => {})[Symbol.dispose]).toBeTypeOf('function') + }) + }) + describe.skipIf(Symbol.dispose)('in environments not supporting it', () => { + it('does not have dispose property', () => { + expect(vi.fn()[Symbol.dispose]).toBeUndefined() + }) + }) + }) + describe('docs example', () => { it('mockClear', () => { const person = {