feat: add support for Explicit Resource Management to mocked functions (#14895)

Author: phryneas

CHANGELOG.md

@@ -18,6 +18,7 @@
 - `[jest-environment-node]` Update jest environment with dispose symbols `Symbol` ([#14888](https://github.com/jestjs/jest/pull/14888) & [#14909](https://github.com/jestjs/jest/pull/14909))
 - `[@jest/fake-timers]` [**BREAKING**] Upgrade `@sinonjs/fake-timers` to v11 ([#14544](https://github.com/jestjs/jest/pull/14544))
 - `[@jest/fake-timers]` Exposing new modern timers function `advanceTimersToFrame()` which advances all timers by the needed milliseconds to execute callbacks currently scheduled with `requestAnimationFrame` ([#14598](https://github.com/jestjs/jest/pull/14598))
+- `[jest-mock]` Add support for the Explicit Resource Management proposal to use the `using` keyword with `jest.spyOn(object, methodName)` ([#14895](https://github.com/jestjs/jest/pull/14895))
 - `[jest-runtime]` Exposing new modern timers function `jest.advanceTimersToFrame()` from `@jest/fake-timers` ([#14598](https://github.com/jestjs/jest/pull/14598))
 - `[jest-runtime]` Support `import.meta.filename` and `import.meta.dirname` (available from [Node 20.11](https://nodejs.org/en/blog/release/v20.11.0))
 - `[@jest/schemas]` Upgrade `@sinclair/typebox` to v0.31 ([#14072](https://github.com/jestjs/jest/pull/14072))

docs/JestObjectAPI.md

@@ -710,6 +710,64 @@ test('plays video', () => {
 });
 ```
 
+#### Spied methods and the `using` keyword
+
+If your codebase is set up to transpile the ["explicit resource management"](https://github.com/tc39/proposal-explicit-resource-management) (e.g. if you are using TypeScript >= 5.2 or the `@babel/plugin-proposal-explicit-resource-management` plugin), you can use `spyOn` in combination with the `using` keyword:
+
+```js
+test('logs a warning', () => {
+  using spy = jest.spyOn(console.warn);
+  doSomeThingWarnWorthy();
+  expect(spy).toHaveBeenCalled();
+});
+```
+
+That code is semantically equal to
+
+```js
+test('logs a warning', () => {
+  let spy;
+  try {
+    spy = jest.spyOn(console.warn);
+    doSomeThingWarnWorthy();
+    expect(spy).toHaveBeenCalled();
+  } finally {
+    spy.mockRestore();
+  }
+});
+```
+
+That way, your spy will automatically be restored to the original value once the current code block is left.
+
+You can even go a step further and use a code block to restrict your mock to only a part of your test without hurting readability.
+
+```js
+test('testing something', () => {
+  {
+    using spy = jest.spyOn(console.warn);
+    setupStepThatWillLogAWarning();
+  }
+  // here, console.warn is already restored to the original value
+  // your test can now continue normally
+});
+```
+
+:::note
+
+If you get a warning that `Symbol.dispose` does not exist, you might need to polyfill that, e.g. with this code:
+
+```js
+if (!Symbol.dispose) {
+  Object.defineProperty(Symbol, 'dispose', {
+    get() {
+      return Symbol.for('nodejs.dispose');
+    },
+  });
+}
+```
+
+:::
+
 ### `jest.spyOn(object, methodName, accessType?)`
 
 Since Jest 22.1.0+, the `jest.spyOn` method takes an optional third argument of `accessType` that can be either `'get'` or `'set'`, which proves to be useful when you want to spy on a getter or a setter, respectively.

e2e/__tests__/explicitResourceManagement.test.ts

@@ -0,0 +1,24 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import {resolve} from 'path';
+import {onNodeVersions} from '@jest/test-utils';
+import {runYarnInstall} from '../Utils';
+import runJest from '../runJest';
+
+const DIR = resolve(__dirname, '../explicit-resource-management');
+
+beforeAll(() => {
+  runYarnInstall(DIR);
+});
+
+onNodeVersions('^18.18.0 || >=20.4.0', () => {
+  test('Explicit resource management is supported', () => {
+    const result = runJest(DIR);
+    expect(result.exitCode).toBe(0);
+  });
+});

e2e/explicit-resource-management/__tests__/index.js

@@ -0,0 +1,47 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+const TestClass = require('../');
+const localClass = new TestClass();
+
+it('restores a mock after a test if it is mocked with a `using` declaration', () => {
+  using mock = jest.spyOn(localClass, 'test').mockImplementation(() => 'ABCD');
+  expect(localClass.test()).toBe('ABCD');
+  expect(localClass.test).toHaveBeenCalledTimes(1);
+  expect(jest.isMockFunction(localClass.test)).toBeTruthy();
+});
+
+it('only sees the unmocked class', () => {
+  expect(localClass.test()).toBe('12345');
+  expect(localClass.test.mock).toBeUndefined();
+  expect(jest.isMockFunction(localClass.test)).toBeFalsy();
+});
+
+test('also works just with scoped code blocks', () => {
+  const scopedInstance = new TestClass();
+  {
+    using mock = jest
+      .spyOn(scopedInstance, 'test')
+      .mockImplementation(() => 'ABCD');
+    expect(scopedInstance.test()).toBe('ABCD');
+    expect(scopedInstance.test).toHaveBeenCalledTimes(1);
+    expect(jest.isMockFunction(scopedInstance.test)).toBeTruthy();
+  }
+  expect(scopedInstance.test()).toBe('12345');
+  expect(scopedInstance.test.mock).toBeUndefined();
+  expect(jest.isMockFunction(scopedInstance.test)).toBeFalsy();
+});
+
+it('jest.fn state should be restored with the `using` keyword', () => {
+  const mock = jest.fn();
+  {
+    using inScope = mock.mockReturnValue(2);
+    expect(inScope()).toBe(2);
+    expect(mock()).toBe(2);
+  }
+  expect(mock()).not.toBe(2);
+});

e2e/explicit-resource-management/babel.config.js

@@ -0,0 +1,10 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+module.exports = {
+  plugins: ['@babel/plugin-proposal-explicit-resource-management'],
+};

e2e/explicit-resource-management/index.js

@@ -0,0 +1,12 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+module.exports = class Test {
+  test() {
+    return '12345';
+  }
+};

e2e/explicit-resource-management/package.json

@@ -0,0 +1,8 @@
+{
+  "jest": {
+    "testEnvironment": "node"
+  },
+  "dependencies": {
+    "@babel/plugin-proposal-explicit-resource-management": "^7.23.9"
+  }
+}

e2e/explicit-resource-management/yarn.lock

@@ -0,0 +1,44 @@
+# This file is generated by running "yarn install" inside your project.
+# Manual changes might be lost - proceed with caution!
+
+__metadata:
+  version: 6
+  cacheKey: 8
+
+"@babel/helper-plugin-utils@npm:^7.22.5":
+  version: 7.22.5
+  resolution: "@babel/helper-plugin-utils@npm:7.22.5"
+  checksum: c0fc7227076b6041acd2f0e818145d2e8c41968cc52fb5ca70eed48e21b8fe6dd88a0a91cbddf4951e33647336eb5ae184747ca706817ca3bef5e9e905151ff5
+  languageName: node
+  linkType: hard
+
+"@babel/plugin-proposal-explicit-resource-management@npm:^7.23.9":
+  version: 7.23.9
+  resolution: "@babel/plugin-proposal-explicit-resource-management@npm:7.23.9"
+  dependencies:
+    "@babel/helper-plugin-utils": ^7.22.5
+    "@babel/plugin-syntax-explicit-resource-management": ^7.23.3
+  peerDependencies:
+    "@babel/core": ^7.0.0-0
+  checksum: d7a37ea28178e251fe289895cf4a37fee47195122a3e172eb088be9b0a55d16d2b2ac3cd6569e9f94c9f9a7744a812f3eba50ec64e3d8f7a48a4e2b0f2caa959
+  languageName: node
+  linkType: hard
+
+"@babel/plugin-syntax-explicit-resource-management@npm:^7.23.3":
+  version: 7.23.3
+  resolution: "@babel/plugin-syntax-explicit-resource-management@npm:7.23.3"
+  dependencies:
+    "@babel/helper-plugin-utils": ^7.22.5
+  peerDependencies:
+    "@babel/core": ^7.0.0-0
+  checksum: 60306808e4680b180a2945d13d4edc7aba91bbd43b300271b89ebd3d3d0bc60f97c6eb7eaa7b9e2f7b61bb0111c24469846f636766517da5385351957c264eb9
+  languageName: node
+  linkType: hard
+
+"root-workspace-0b6124@workspace:.":
+  version: 0.0.0-use.local
+  resolution: "root-workspace-0b6124@workspace:."
+  dependencies:
+    "@babel/plugin-proposal-explicit-resource-management": ^7.23.9
+  languageName: unknown
+  linkType: soft

packages/jest-environment-node/src/index.ts

@@ -155,9 +155,9 @@ export default class NodeEnvironment implements JestEnvironment<Timer> {
     if ('asyncDispose' in Symbol && !('asyncDispose' in global.Symbol)) {
       const globalSymbol = global.Symbol as unknown as SymbolConstructor;
       // @ts-expect-error - it's readonly - but we have checked above that it's not there
-      globalSymbol.asyncDispose = globalSymbol('nodejs.asyncDispose');
+      globalSymbol.asyncDispose = globalSymbol.for('nodejs.asyncDispose');
       // @ts-expect-error - it's readonly - but we have checked above that it's not there
-      globalSymbol.dispose = globalSymbol('nodejs.dispose');
+      globalSymbol.dispose = globalSymbol.for('nodejs.dispose');
     }
 
     // Node's error-message stack size is limited at 10, but it's pretty useful

packages/jest-mock/src/index.ts

@@ -5,6 +5,8 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+/// <reference lib="ESNext.Disposable" />
+
 /* eslint-disable local/ban-types-eventually, local/prefer-rest-params-eventually */
 
 import {isPromise} from 'jest-util';
@@ -131,7 +133,8 @@ type ResolveType<T extends FunctionLike> =
 type RejectType<T extends FunctionLike> =
   ReturnType<T> extends PromiseLike<any> ? unknown : never;
 
-export interface MockInstance<T extends FunctionLike = UnknownFunction> {
+export interface MockInstance<T extends FunctionLike = UnknownFunction>
+  extends Disposable {
   _isMockFunction: true;
   _protoImpl: Function;
   getMockImplementation(): T | undefined;
@@ -797,6 +800,9 @@ export class ModuleMocker {
       };
 
       f.withImplementation = withImplementation.bind(this);
+      if (Symbol.dispose) {
+        f[Symbol.dispose] = f.mockRestore;
+      }
 
       function withImplementation(fn: T, callback: () => void): void;
       function withImplementation(