Fix isolateModules not working when module has already been imported

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

@@ -240,6 +240,105 @@ describe('isolateModules', () => {
       expect(exports.getState()).toBe(1);
     }));
 
+  it('isolates module after module has already been required', () =>
+    createRuntime(__filename, {
+      moduleNameMapper,
+    }).then(runtime => {
+      let exports;
+
+      exports = runtime.requireModuleOrMock(
+        runtime.__mockRootPath,
+        'ModuleWithState',
+      );
+      exports.increment();
+      expect(exports.getState()).toBe(2);
+
+      runtime.isolateModules(() => {
+        exports = runtime.requireModuleOrMock(
+          runtime.__mockRootPath,
+          'ModuleWithState',
+        );
+        expect(exports.getState()).toBe(1);
+      });
+    }));
+
+  it('can isolate the same module multiple times', () =>
+    createRuntime(__filename, {
+      moduleNameMapper,
+    }).then(runtime => {
+      runtime.isolateModules(() => {
+        const exports = runtime.requireModuleOrMock(
+          runtime.__mockRootPath,
+          'ModuleWithState',
+        );
+        exports.increment();
+        expect(exports.getState()).toBe(2);
+      });
+
+      runtime.isolateModules(() => {
+        const exports = runtime.requireModuleOrMock(
+          runtime.__mockRootPath,
+          'ModuleWithState',
+        );
+        expect(exports.getState()).toBe(1);
+      });
+    }));
+
+  it("doesn't isolate higher dependencies within the isolated module", () =>
+    createRuntime(__filename, {
+      moduleNameMapper,
+    }).then(runtime => {
+      const exports = runtime.requireModuleOrMock(
+        runtime.__mockRootPath,
+        'ModuleWithStatefulImport',
+      );
+      exports.increment();
+      expect(exports.getState()).toBe(2);
+
+      let isolatedExports;
+      runtime.isolateModules(
+        () =>
+          (isolatedExports = runtime.requireModuleOrMock(
+            runtime.__mockRootPath,
+            'ModuleWithStatefulImport',
+          )),
+      );
+
+      // Higher dependency's state is unaffected
+      expect(exports.getState()).toBe(isolatedExports.getState());
+      // Internal object within isolated module becomes a copy
+      expect(exports.scopedObject).toEqual(isolatedExports.scopedObject);
+      expect(exports.scopedObject).not.toBe(isolatedExports.scopedObject);
+    }));
+
+  it('can isolate mocks', () =>
+    createRuntime(__filename, {
+      moduleNameMapper,
+    }).then(runtime => {
+      runtime.setMock(runtime.__mockRootPath, 'ModuleWithState', () => {
+        let state = 50;
+        return {
+          getState: () => state,
+          increment: () => state++,
+        };
+      });
+
+      const exports = runtime.requireModuleOrMock(
+        runtime.__mockRootPath,
+        'ModuleWithStatefulImport',
+      );
+      exports.increment();
+      expect(exports.getState()).toBe(51);
+
+      runtime.isolateModules(() => {
+        const isolatedMock = runtime.requireMock(
+          runtime.__mockRootPath,
+          'ModuleWithState',
+        );
+        expect(isolatedMock.getState()).toBe(50);
+      });
+    }));
+
   it('cannot nest isolateModules blocks', () =>
     createRuntime(__filename, {
       moduleNameMapper,

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

@@ -0,0 +1,5 @@
+import {getState, increment} from './ModuleWithState';
+
+export {getState, increment};
+
+export const scopedObject = {};

packages/jest-runtime/src/index.ts

@@ -90,6 +90,7 @@ class Runtime {
   private _explicitShouldMock: BooleanObject;
   private _internalModuleRegistry: ModuleRegistry;
   private _isCurrentlyExecutingManualMock: string | null;
+  private _isolatedRequireInProgress: boolean;
   private _mockFactories: Record<string, () => unknown>;
   private _mockMetaDataCache: {
     [key: string]: MockFunctionMetadata<unknown, Array<unknown>>;
@@ -130,6 +131,7 @@ class Runtime {
     this._explicitShouldMock = Object.create(null);
     this._internalModuleRegistry = new Map();
     this._isCurrentlyExecutingManualMock = null;
+    this._isolatedRequireInProgress = false;
     this._mockFactories = Object.create(null);
     this._mockRegistry = new Map();
     // during setup, this cannot be null (and it's fine to explode if it is)
@@ -318,48 +320,52 @@ class Runtime {
       modulePath = this._resolveModule(from, moduleName);
     }
 
-    let moduleRegistry;
+    let disableIsolatedRequireInProgress = false;
 
-    if (!options || !options.isInternalModule) {
-      if (
-        this._moduleRegistry.get(modulePath) ||
-        !this._isolatedModuleRegistry
-      ) {
-        moduleRegistry = this._moduleRegistry;
-      } else {
-        moduleRegistry = this._isolatedModuleRegistry;
-      }
-    } else {
+    let moduleRegistry = this._moduleRegistry;
+    if (options && options.isInternalModule) {
       moduleRegistry = this._internalModuleRegistry;
+    } else if (
+      this._isolatedModuleRegistry &&
+      !this._isolatedRequireInProgress
+    ) {
+      // Only isolate modules that were explicitly required inside isolateModules callback,
+      // preventing higher dependencies within that module to become indirectly isolated
+      this._isolatedRequireInProgress = true;
+      disableIsolatedRequireInProgress = true;
+      moduleRegistry = this._isolatedModuleRegistry
     }
 
-    const module = moduleRegistry.get(modulePath);
-    if (module) {
-      return module.exports;
-    }
-
-    // We must register the pre-allocated module object first so that any
-    // circular dependencies that may arise while evaluating the module can
-    // be satisfied.
-    const localModule: InitialModule = {
-      children: [],
-      exports: {},
-      filename: modulePath,
-      id: modulePath,
-      loaded: false,
-    };
-    moduleRegistry.set(modulePath, localModule);
+    let module = moduleRegistry.get(modulePath);
+    if (!module) {
+      // We must register the pre-allocated module object first so that any
+      // circular dependencies that may arise while evaluating the module can
+      // be satisfied.
+      const localModule: InitialModule = {
+        children: [],
+        exports: {},
+        filename: modulePath,
+        id: modulePath,
+        loaded: false,
+      };
+      moduleRegistry.set(modulePath, localModule);
+
+      this._loadModule(
+        localModule,
+        from,
+        moduleName,
+        modulePath,
+        options,
+        moduleRegistry,
+      );
 
-    this._loadModule(
-      localModule,
-      from,
-      moduleName,
-      modulePath,
-      options,
-      moduleRegistry,
-    );
+      module = localModule;
+    }
 
-    return localModule.exports;
+    if (disableIsolatedRequireInProgress) {
+      this._isolatedRequireInProgress = false;
+    }
+    return module.exports;
   }
 
   requireInternalModule(from: Config.Path, to?: string) {
@@ -377,20 +383,25 @@ class Runtime {
       moduleName,
     );
 
-    if (
-      this._isolatedMockRegistry &&
-      this._isolatedMockRegistry.get(moduleID)
-    ) {
-      return this._isolatedMockRegistry.get(moduleID);
-    } else if (this._mockRegistry.get(moduleID)) {
-      return this._mockRegistry.get(moduleID);
-    }
+    let disableIsolatedRequireInProgress = false;
 
-    const mockRegistry = this._isolatedMockRegistry || this._mockRegistry;
+    let mockRegistry = this._mockRegistry;
+    if (this._isolatedMockRegistry && !this._isolatedRequireInProgress) {
+      this._isolatedRequireInProgress = true;
+      disableIsolatedRequireInProgress = true;
+      mockRegistry = this._isolatedMockRegistry;
+    }
 
-    if (moduleID in this._mockFactories) {
-      const module = this._mockFactories[moduleID]();
+    let module = mockRegistry.get(moduleID);
+    if (!module && moduleID in this._mockFactories) {
+      module = this._mockFactories[moduleID]();
       mockRegistry.set(moduleID, module);
+    }
+
+    if (module) {
+      if (disableIsolatedRequireInProgress) {
+        this._isolatedRequireInProgress = false;
+      }
       return module;
     }
 
@@ -454,6 +465,9 @@ class Runtime {
       mockRegistry.set(moduleID, this._generateMock(from, moduleName));
     }
 
+    if (disableIsolatedRequireInProgress) {
+      this._isolatedRequireInProgress = false;
+    }
     return mockRegistry.get(moduleID);
   }