feat(core): support using for mock restore (#1293)

Author: 9aoy

e2e/spy/dispose.test.ts

@@ -0,0 +1,29 @@
+import { describe, expect, it, rstest } from '@rstest/core';
+
+describe('mock dispose', () => {
+  it('restores spies when leaving a using scope', () => {
+    const hi = {
+      sayHi: () => 'hi',
+    };
+
+    {
+      using spy = rstest.spyOn(hi, 'sayHi').mockImplementation(() => 'hello');
+
+      expect(hi.sayHi()).toBe('hello');
+      expect(spy).toHaveBeenCalledTimes(1);
+    }
+
+    expect(hi.sayHi()).toBe('hi');
+    expect(rstest.isMockFunction(hi.sayHi)).toBeFalsy();
+  });
+
+  it('resets mock functions when disposed manually', () => {
+    const sayHi = rstest.fn(() => 'hi').mockImplementation(() => 'hello');
+
+    expect(sayHi()).toBe('hello');
+
+    sayHi[Symbol.dispose]();
+
+    expect(sayHi()).toBe('hi');
+  });
+});

packages/core/src/runtime/api/spy.ts

@@ -204,6 +204,15 @@ export const initSpy = (): Pick<
       mockName = mockFn?.name;
     };
 
+    if (Symbol.dispose) {
+      Object.defineProperty(spyFn, Symbol.dispose, {
+        value: () => {
+          spyFn.mockRestore();
+        },
+        configurable: true,
+      });
+    }
+
     mocks.add(spyFn);
 
     return spyFn;

packages/core/src/types/mock.ts

@@ -116,6 +116,10 @@ export interface MockInstance<T extends FunctionLike = FunctionLike> {
    * Does what `mockReset` does and restores original descriptors of spied-on objects.
    */
   mockRestore(): void;
+  /**
+   * Restores the mock when it leaves a `using` scope.
+   */
+  [Symbol.dispose](): void;
   /**
    * Returns current mock implementation if there is one.
    */

website/docs/en/api/runtime-api/rstest/mock-instance.mdx

@@ -4,6 +4,8 @@ description: MockInstance is the type of all mock/spy instances, providing a ric
 overviewHeaders: [2, 3]
 ---
 
+import { ApiMeta } from '@components/ApiMeta';
+
 # Mock instance
 
 `MockInstance` is the type of all mock/spy instances, providing a rich set of APIs for controlling and inspecting mocks.
@@ -58,7 +60,7 @@ fn.mockReset();
 
 ## mockRestore
 
-- **Type:** `() => MockInstance`
+- **Type:** `() => void`
 
 Restores the original method of a spied object (only effective for spies).
 
@@ -68,6 +70,25 @@ const spy = rstest.spyOn(obj, 'foo');
 spy.mockRestore();
 ```
 
+## Symbol.dispose
+
+<ApiMeta addedVersion="0.10.2" />
+
+- **Type:** `() => void`
+
+Calls `mockRestore()` when the mock leaves a `using` scope.
+
+```ts
+const obj = { foo: () => 1 };
+
+{
+  using spy = rstest.spyOn(obj, 'foo').mockImplementation(() => 2);
+  expect(obj.foo()).toBe(2);
+}
+
+expect(obj.foo()).toBe(1);
+```
+
 ## getMockImplementation
 
 - **Type:** `() => Function | undefined`

website/docs/en/guide/basic/mock.mdx

@@ -2,6 +2,8 @@
 description: Mock functions, object methods, ESM modules, and CommonJS modules in Rstest, and distinguish mock state from module state.
 ---
 
+import { ApiMeta } from '@components/ApiMeta';
+
 # Mocking
 
 Mocking lets you replace dependencies in tests, control return values, and assert how functions or modules are called. Rstest provides different mocking APIs for functions, object methods, ESM modules, CommonJS modules, and object trees.
@@ -211,6 +213,30 @@ test('logs a warning when validation fails', () => {
 });
 ```
 
+### `using` syntax
+
+<ApiMeta addedVersion="0.10.2" />
+
+Rstest supports the `using` syntax to restore the spy automatically when the block exits:
+
+```ts title="logger.test.ts"
+import { expect, rstest, test } from '@rstest/core';
+
+test('logs a warning when validation fails', () => {
+  {
+    using warn = rstest
+      .spyOn(console, 'warn')
+      .mockImplementation(() => undefined);
+
+    console.warn('invalid payload');
+
+    expect(warn).toHaveBeenCalledWith('invalid payload');
+  }
+
+  // console.warn is restored here
+});
+```
+
 This pattern is commonly used with globals such as `console` and `Date`, as well as shared objects that already exist in the test.
 
 ## Deep-mock objects

website/docs/zh/api/runtime-api/rstest/mock-instance.mdx

@@ -4,6 +4,8 @@ description: MockInstance 是所有 mock 和 spy 实例的类型。
 overviewHeaders: [2, 3]
 ---
 
+import { ApiMeta } from '@components/ApiMeta';
+
 # Mock instance
 
 `MockInstance` 是所有 mock 和 spy 实例的类型。
@@ -58,7 +60,7 @@ fn.mockReset();
 
 ## mockRestore
 
-- **类型：** `() => MockInstance`
+- **类型：** `() => void`
 
 恢复被 spy 的对象的原始方法（仅对 spy 有效）。
 
@@ -68,6 +70,25 @@ const spy = rstest.spyOn(obj, 'foo');
 spy.mockRestore();
 ```
 
+## Symbol.dispose
+
+<ApiMeta addedVersion="0.10.2" />
+
+- **类型：** `() => void`
+
+当 mock 离开 `using` 作用域时调用 `mockRestore()`。
+
+```ts
+const obj = { foo: () => 1 };
+
+{
+  using spy = rstest.spyOn(obj, 'foo').mockImplementation(() => 2);
+  expect(obj.foo()).toBe(2);
+}
+
+expect(obj.foo()).toBe(1);
+```
+
 ## getMockImplementation
 
 - **类型：** `() => Function | undefined`

website/docs/zh/guide/basic/mock.mdx

@@ -2,6 +2,8 @@
 description: 使用 Rstest mock 函数、对象方法、ESM 模块、CommonJS 模块，并区分 mock 状态和模块状态。
 ---
 
+import { ApiMeta } from '@components/ApiMeta';
+
 # Mocking
 
 Mocking 用于在测试中替换依赖、控制返回值，并断言函数或模块的调用行为。Rstest 提供了多组 mock API，分别适用于函数、对象方法、ESM 模块、CommonJS 模块和对象树。
@@ -211,6 +213,30 @@ test('logs a warning when validation fails', () => {
 });
 ```
 
+### `using` 语法
+
+<ApiMeta addedVersion="0.10.2" />
+
+Rstest 支持使用 `using` 语法在代码块退出时自动恢复 spy：
+
+```ts title="logger.test.ts"
+import { expect, rstest, test } from '@rstest/core';
+
+test('logs a warning when validation fails', () => {
+  {
+    using warn = rstest
+      .spyOn(console, 'warn')
+      .mockImplementation(() => undefined);
+
+    console.warn('invalid payload');
+
+    expect(warn).toHaveBeenCalledWith('invalid payload');
+  }
+
+  // console.warn 在这里已恢复
+});
+```
+
 这类写法常用于 `console`、`Date` 这类全局对象，以及测试里已经存在的共享对象。
 
 ## 深度 mock 对象