[Feature]: add isolateModulesAsync

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ### Features
 
+- `[jest-runtime]` Add `jest.isolateModulesAsync` for scoped module initialization of asynchronous functions ([#13680](https://github.com/facebook/jest/pull/13680))
+
 ### Fixes
 
 - `[jest-resolve]` add global paths to `require.resolve.paths` ([#13633](https://github.com/facebook/jest/pull/13633))

docs/JestObjectAPI.md

@@ -568,6 +568,32 @@ jest.isolateModules(() => {
 const otherCopyOfMyModule = require('myModule');
 ```
 
+### `jest.isolateModulesAsync(fn)`
+
+`jest.isolateModulesAsync()` is the equivalent of `jest.isolateModules()`, but for async functions that need to be wrapped. The caller is expected to `await` the completion of `isolateModulesAsync`.
+
+```js
+let myModule;
+await jest.isolateModulesAsync(async () => {
+  myModule = await import('myModule');
+  // do async stuff here
+});
+
+const otherCopyOfMyModule = require('myModule');
+```
+
+`jest.isolateModulesAsync()` can also run synchronous callbacks:
+
+```js
+let myModule;
+await jest.isolateModulesAsync(() => {
+  myModule = require('myModule');
+  // do sync stuff here
+});
+
+const otherCopyOfMyModule = require('myModule');
+```
+
 ## Mock Functions
 
 ### `jest.fn(implementation?)`

packages/jest-environment/src/index.ts

@@ -172,6 +172,12 @@ export interface Jest {
    * local module state doesn't conflict between tests.
    */
   isolateModules(fn: () => void): Jest;
+  /**
+   * `jest.isolateModulesAsync()` is the equivalent of `jest.isolateModules()`, but for
+   * async functions to be wrapped. The caller is expected to `await` the completion of
+   * `isolateModulesAsync`.
+   */
+  isolateModulesAsync(fn: () => Promise<void>): Promise<void>;
   /**
    * Mocks a module with an auto-mocked version when it is being required.
    */

packages/jest-runtime/src/__tests__/runtime_require_module_or_mock.test.js

@@ -214,7 +214,7 @@ describe('resetModules', () => {
 });
 
 describe('isolateModules', () => {
-  it("keeps it's registry isolated from global one", async () => {
+  it('keeps its registry isolated from global one', async () => {
     const runtime = await createRuntime(__filename, {
       moduleNameMapper,
     });
@@ -287,7 +287,7 @@ describe('isolateModules', () => {
         runtime.isolateModules(() => {});
       });
     }).toThrow(
-      'isolateModules cannot be nested inside another isolateModules.',
+      'isolateModules cannot be nested inside another isolateModules or isolateModulesAsync.',
     );
   });
 
@@ -325,6 +325,7 @@ describe('isolateModules', () => {
     beforeEach(() => {
       jest.isolateModules(() => {
         exports = require('./test_root/ModuleWithState');
+        exports.set(1); // Ensure idempotency with the isolateModulesAsync test
       });
     });
 
@@ -340,3 +341,242 @@ describe('isolateModules', () => {
     });
   });
 });
+
+describe('isolateModulesAsync', () => {
+  describe('with sync callback', () => {
+    it('keeps its registry isolated from global one', async () => {
+      const runtime = await createRuntime(__filename, {
+        moduleNameMapper,
+      });
+      let exports;
+      exports = runtime.requireModuleOrMock(
+        runtime.__mockRootPath,
+        'ModuleWithState',
+      );
+      exports.increment();
+      expect(exports.getState()).toBe(2);
+
+      await runtime.isolateModulesAsync(() => {
+        exports = runtime.requireModuleOrMock(
+          runtime.__mockRootPath,
+          'ModuleWithState',
+        );
+        expect(exports.getState()).toBe(1);
+      });
+
+      exports = runtime.requireModuleOrMock(
+        runtime.__mockRootPath,
+        'ModuleWithState',
+      );
+      expect(exports.getState()).toBe(2);
+    });
+
+    it('resets all modules after the block', async () => {
+      const runtime = await createRuntime(__filename, {
+        moduleNameMapper,
+      });
+      let exports;
+      await runtime.isolateModulesAsync(() => {
+        exports = runtime.requireModuleOrMock(
+          runtime.__mockRootPath,
+          'ModuleWithState',
+        );
+        expect(exports.getState()).toBe(1);
+        exports.increment();
+        expect(exports.getState()).toBe(2);
+      });
+
+      exports = runtime.requireModuleOrMock(
+        runtime.__mockRootPath,
+        'ModuleWithState',
+      );
+      expect(exports.getState()).toBe(1);
+    });
+
+    it('resets module after failing', async () => {
+      const runtime = await createRuntime(__filename, {
+        moduleNameMapper,
+      });
+      await expect(
+        runtime.isolateModulesAsync(() => {
+          throw new Error('Error from isolated module');
+        }),
+      ).rejects.toThrow('Error from isolated module');
+
+      await runtime.isolateModulesAsync(() => {
+        expect(true).toBe(true);
+      });
+    });
+
+    it('cannot nest isolateModulesAsync blocks', async () => {
+      const runtime = await createRuntime(__filename, {
+        moduleNameMapper,
+      });
+      await expect(async () => {
+        await runtime.isolateModulesAsync(async () => {
+          await runtime.isolateModulesAsync(() => {});
+        });
+      }).rejects.toThrow(
+        'isolateModulesAsync cannot be nested inside another isolateModulesAsync or isolateModules.',
+      );
+    });
+
+    it('can call resetModules within a isolateModules block', async () => {
+      const runtime = await createRuntime(__filename, {
+        moduleNameMapper,
+      });
+      let exports;
+      await runtime.isolateModulesAsync(() => {
+        exports = runtime.requireModuleOrMock(
+          runtime.__mockRootPath,
+          'ModuleWithState',
+        );
+        expect(exports.getState()).toBe(1);
+
+        exports.increment();
+        runtime.resetModules();
+
+        exports = runtime.requireModuleOrMock(
+          runtime.__mockRootPath,
+          'ModuleWithState',
+        );
+        expect(exports.getState()).toBe(1);
+      });
+
+      exports = runtime.requireModuleOrMock(
+        runtime.__mockRootPath,
+        'ModuleWithState',
+      );
+      expect(exports.getState()).toBe(1);
+    });
+  });
+
+  describe('with async callback', () => {
+    it('keeps its registry isolated from global one', async () => {
+      const runtime = await createRuntime(__filename, {
+        moduleNameMapper,
+      });
+      let exports;
+      exports = runtime.requireModuleOrMock(
+        runtime.__mockRootPath,
+        'ModuleWithState',
+      );
+      exports.increment();
+      expect(exports.getState()).toBe(2);
+
+      await runtime.isolateModulesAsync(async () => {
+        exports = runtime.requireModuleOrMock(
+          runtime.__mockRootPath,
+          'ModuleWithState',
+        );
+        expect(exports.getState()).toBe(1);
+      });
+
+      exports = runtime.requireModuleOrMock(
+        runtime.__mockRootPath,
+        'ModuleWithState',
+      );
+      expect(exports.getState()).toBe(2);
+    });
+
+    it('resets all modules after the block', async () => {
+      const runtime = await createRuntime(__filename, {
+        moduleNameMapper,
+      });
+      let exports;
+      await runtime.isolateModulesAsync(async () => {
+        exports = runtime.requireModuleOrMock(
+          runtime.__mockRootPath,
+          'ModuleWithState',
+        );
+        expect(exports.getState()).toBe(1);
+        exports.increment();
+        expect(exports.getState()).toBe(2);
+      });
+
+      exports = runtime.requireModuleOrMock(
+        runtime.__mockRootPath,
+        'ModuleWithState',
+      );
+      expect(exports.getState()).toBe(1);
+    });
+
+    it('resets module after failing', async () => {
+      const runtime = await createRuntime(__filename, {
+        moduleNameMapper,
+      });
+      await expect(
+        runtime.isolateModulesAsync(async () => {
+          throw new Error('Error from isolated module');
+        }),
+      ).rejects.toThrow('Error from isolated module');
+
+      await runtime.isolateModulesAsync(async () => {
+        expect(true).toBe(true);
+      });
+    });
+
+    it('cannot nest isolateModulesAsync blocks', async () => {
+      const runtime = await createRuntime(__filename, {
+        moduleNameMapper,
+      });
+      await expect(async () => {
+        await runtime.isolateModulesAsync(async () => {
+          await runtime.isolateModulesAsync(() => Promise.resolve());
+        });
+      }).rejects.toThrow(
+        'isolateModulesAsync cannot be nested inside another isolateModulesAsync or isolateModules.',
+      );
+    });
+
+    it('can call resetModules within a isolateModules block', async () => {
+      const runtime = await createRuntime(__filename, {
+        moduleNameMapper,
+      });
+      let exports;
+      await runtime.isolateModulesAsync(async () => {
+        exports = runtime.requireModuleOrMock(
+          runtime.__mockRootPath,
+          'ModuleWithState',
+        );
+        expect(exports.getState()).toBe(1);
+
+        exports.increment();
+        runtime.resetModules();
+
+        exports = runtime.requireModuleOrMock(
+          runtime.__mockRootPath,
+          'ModuleWithState',
+        );
+        expect(exports.getState()).toBe(1);
+      });
+
+      exports = runtime.requireModuleOrMock(
+        runtime.__mockRootPath,
+        'ModuleWithState',
+      );
+      expect(exports.getState()).toBe(1);
+    });
+  });
+
+  describe('can use isolateModulesAsync from a beforeEach block', () => {
+    let exports;
+    beforeEach(async () => {
+      await jest.isolateModulesAsync(async () => {
+        exports = require('./test_root/ModuleWithState');
+        exports.set(1); // Ensure idempotency with the isolateModules test
+      });
+    });
+
+    it('can use the required module from beforeEach and re-require it', () => {
+      expect(exports.getState()).toBe(1);
+      exports.increment();
+      expect(exports.getState()).toBe(2);
+
+      exports = require('./test_root/ModuleWithState');
+      expect(exports.getState()).toBe(2);
+      exports.increment();
+      expect(exports.getState()).toBe(3);
+    });
+  });
+});

packages/jest-runtime/src/__tests__/test_root/ModuleWithState.js

@@ -8,6 +8,10 @@
 
 let state = 1;
 
+export const set = i => {
+  state = i;
+};
+
 export const increment = () => {
   state += 1;
 };

packages/jest-runtime/src/index.ts

@@ -1125,7 +1125,7 @@ export default class Runtime {
   isolateModules(fn: () => void): void {
     if (this._isolatedModuleRegistry || this._isolatedMockRegistry) {
       throw new Error(
-        'isolateModules cannot be nested inside another isolateModules.',
+        'isolateModules cannot be nested inside another isolateModules or isolateModulesAsync.',
       );
     }
     this._isolatedModuleRegistry = new Map();
@@ -1141,6 +1141,25 @@ export default class Runtime {
     }
   }
 
+  async isolateModulesAsync(fn: () => Promise<void> | void): Promise<void> {
+    if (this._isolatedModuleRegistry || this._isolatedMockRegistry) {
+      throw new Error(
+        'isolateModulesAsync cannot be nested inside another isolateModulesAsync or isolateModules.',
+      );
+    }
+    this._isolatedModuleRegistry = new Map();
+    this._isolatedMockRegistry = new Map();
+    try {
+      await fn();
+    } finally {
+      // might be cleared within the callback
+      this._isolatedModuleRegistry?.clear();
+      this._isolatedMockRegistry?.clear();
+      this._isolatedModuleRegistry = null;
+      this._isolatedMockRegistry = null;
+    }
+  }
+
   resetModules(): void {
     this._isolatedModuleRegistry?.clear();
     this._isolatedMockRegistry?.clear();
@@ -2226,6 +2245,7 @@ export default class Runtime {
       getTimerCount: () => _getFakeTimers().getTimerCount(),
       isMockFunction: this._moduleMocker.isMockFunction,
       isolateModules,
+      isolateModulesAsync: this.isolateModulesAsync,
       mock,
       mocked,
       now: () => _getFakeTimers().now(),

packages/jest-types/__typetests__/jest.test.ts

@@ -98,6 +98,10 @@ expectError(jest.enableAutomock('moduleName'));
 expectType<typeof jest>(jest.isolateModules(() => {}));
 expectError(jest.isolateModules());
 
+expectType<Promise<void>>(jest.isolateModulesAsync(async () => {}));
+expectType<Promise<void>>(jest.isolateModulesAsync(() => {}));
+expectError(jest.isolateModulesAsync());
+
 expectType<typeof jest>(jest.mock('moduleName'));
 expectType<typeof jest>(jest.mock('moduleName', jest.fn()));
 expectType<typeof jest>(