feat: restore spies automatically with using

optional-types.d.ts

@@ -0,0 +1,4 @@
+export interface Disposable {
+  // @ts-ignore -- Symbol.dispose might not be in user types
+  [Symbol.dispose]: () => void
+}

package.json

@@ -29,7 +29,8 @@
   "module": "./dist/index.js",
   "types": "./dist/index.d.ts",
   "files": [
-    "dist"
+    "dist",
+    "optional-types.d.ts"
   ],
   "engines": {
     "node": ">=14.0.0"

src/internal.ts

@@ -1,5 +1,6 @@
 import { assert, define, defineValue, isPromise, isType } from './utils.js'
 import { SYMBOL_STATE } from './constants.js'
+import type { Disposable } from '../optional-types.js'
 
 interface GetState {
   <A extends any[], R>(spy: SpyInternalImpl<A, R>): SpyInternalImplState<A, R>
@@ -67,7 +68,9 @@ export interface Spy<A extends any[] = any[], R = any>
   nextResult(result: R): this
 }
 
-export interface SpyImpl<A extends any[] = any[], R = any> extends Spy<A, R> {
+export interface SpyImpl<A extends any[] = any[], R = any>
+  extends Spy<A, R>,
+    Disposable {
   getOriginal(): (...args: A) => R
   willCall(cb: (...args: A) => R): this
   restore(): void

src/spyOn.ts

@@ -261,6 +261,9 @@ export function spyOn<T extends object, K extends string & keyof T>(
   ;(['restore', 'getOriginal', 'willCall'] as const).forEach((method) => {
     defineValue(spy, method, spy[SYMBOL_STATE][method])
   })
+  if (Symbol.dispose) {
+    defineValue(spy, Symbol.dispose, spy[SYMBOL_STATE].restore)
+  }
   return spy as any as SpyImpl
 }
 

test/index.test.ts

@@ -82,6 +82,39 @@ test('can spy on method', () => {
   expect(method.results).toEqual([ok('a!'), ok('b!'), ok('C!'), error(err)])
 })
 
+test('restores mock on leaving a using scope', () => {
+  const obj = {
+    method: () => 'original',
+  }
+  const original = obj.method
+
+  {
+    using method = spyOn(obj, 'method').willCall(() => 'mocked')
+
+    expect(obj.method()).toBe('mocked')
+    expect(method.called).toBe(true)
+  }
+
+  expect(obj.method).toBe(original)
+})
+
+test('restores mock when a using scope throws', () => {
+  const obj = {
+    method: () => 'original',
+  }
+  const original = obj.method
+
+  expect(() => {
+    using method = spyOn(obj, 'method').willCall(() => 'mocked')
+
+    expect(obj.method()).toBe('mocked')
+    expect(method.called).toBe(true)
+    throw new Error('test')
+  }).toThrowError('test')
+
+  expect(obj.method).toBe(original)
+})
+
 test('resets all spies', () => {
   let one = {
     method(arg: string) {

tsup.config.ts

@@ -2,6 +2,7 @@ import { defineConfig } from 'tsup'
 
 export default defineConfig({
   entryPoints: ['src/index.ts'],
+  external: ['../optional-types.js'],
   outDir: 'dist',
   format: ['esm'],
   tsconfig: './tsconfig.json',