make doMock() return a disposable

Author: kirkwaiblinger

docs/api/vi.md

@@ -179,11 +179,11 @@ If there is no `__mocks__` folder or a factory provided, Vitest will import the
 function doMock(
   path: string,
   factory?: MockOptions | MockFactory<unknown>
-): void
+): Disposable
 function doMock<T>(
   module: Promise<T>,
   factory?: MockOptions | MockFactory<T>
-): void
+): Disposable
 ```
 
 The same as [`vi.mock`](#vi-mock), but it's not hoisted to the top of the file, so you can reference variables in the global file scope. The next [dynamic import](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Operators/import) of the module will be mocked.
@@ -229,6 +229,26 @@ test('importing the next module imports mocked one', async () => {
 })
 ```
 
+::: 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 [`vi.doUnmock()`](#vi-dounmock) on the mocked module when the containing block is exited. This is especially useful when mocking a dynamically imported module for a single test case.
+
+```ts
+it('uses a mocked version of my-module', () => {
+  using _mockDisposable = vi.doMock('my-module')
+
+  const myModule = await import('my-module') // mocked
+
+  // my-module is restored here
+})
+
+it('uses the normal version of my-module again', () => {
+  const myModule = await import('my-module') // not mocked
+})
+```
+:::
+
+::: tip
+
 ### vi.mocked
 
 ```ts

packages/vitest/src/integrations/vi.ts

@@ -220,11 +220,13 @@ export interface VitestUtils {
    * Mocking algorithm is described in [documentation](https://vitest.dev/guide/mocking/modules).
    * @param path Path to the module. Can be aliased, if your Vitest config supports it
    * @param factory Mocked module factory. The result of this function will be an exports object
+   *
+   * @returns A disposable object that calls {@link doUnmock()} when disposed
    */
   // eslint-disable-next-line ts/method-signature-style
-  doMock(path: string, factory?: MockFactoryWithHelper | MockOptions): void
+  doMock(path: string, factory?: MockFactoryWithHelper | MockOptions): Disposable
   // eslint-disable-next-line ts/method-signature-style
-  doMock<T>(module: Promise<T>, factory?: MockFactoryWithHelper<T> | MockOptions): void
+  doMock<T>(module: Promise<T>, factory?: MockFactoryWithHelper<T> | MockOptions): Disposable
   /**
    * Removes module from mocked registry. All subsequent calls to import will return original module.
    *
@@ -617,6 +619,14 @@ function createVitest(): VitestUtils {
               )
           : factory,
       )
+
+      const rv = {} as Disposable
+      if (Symbol.dispose) {
+        rv[Symbol.dispose] = () => {
+          _mocker().queueUnmock(path, importer)
+        }
+      }
+      return rv
     },
 
     doUnmock(path: string | Promise<unknown>) {

test/core/test/do-mock.test.ts

@@ -30,3 +30,27 @@ test('the second doMock can override the first doMock', async () => {
 
   expect(incrementWith20(1)).toBe(21)
 })
+
+test('doMock works with using', async () => {
+  vi.doUnmock('./fixtures/increment')
+
+  {
+    const { increment: incrementWith1 } = await import('./fixtures/increment')
+    expect(incrementWith1(1)).toBe(2)
+  }
+
+  {
+    using _incrementMock = vi.doMock('./fixtures/increment', () => ({
+      increment: (num: number) => num + 10,
+    }))
+
+    const { increment: incrementWith10 } = await import('./fixtures/increment')
+
+    expect(incrementWith10(1)).toBe(11)
+  }
+
+  {
+    const { increment: incrementWith1 } = await import('./fixtures/increment')
+    expect(incrementWith1(1)).toBe(2)
+  }
+})