Fix isolateModules not working correctly when module has already been imported

CHANGELOG.md

@@ -7,6 +7,7 @@
 ### Fixes
 
 - `[jest-core]` Capture execError during `TestScheduler.scheduleTests` and dispatch to reporters ([#13203](https://github.com/facebook/jest/pull/13203))
+- `[jest-runtime]` Fix `isolateModules` not working correctly when module has already been imported ([#8634](https://github.com/facebook/jest/pull/8634))
 
 ### Chore & Maintenance
 

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

@@ -278,6 +278,105 @@ describe('isolateModules', () => {
     });
   });
 
+  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', async () => {
     const runtime = await createRuntime(__filename, {
       moduleNameMapper,

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

@@ -0,0 +1,13 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ */
+
+import {getState, increment} from './ModuleWithState';
+
+export {getState, increment};
+
+export const scopedObject = {};

packages/jest-runtime/src/index.ts

@@ -166,6 +166,7 @@ export default class Runtime {
     | null;
   private readonly _internalModuleRegistry: ModuleRegistry;
   private _isCurrentlyExecutingManualMock: string | null;
+  private _isolatedRequireInProgress: boolean;
   private _mainModule: Module | null;
   private readonly _mockFactories: Map<string, () => unknown>;
   private readonly _mockMetaDataCache: Map<string, MockMetadata<any>>;
@@ -224,6 +225,7 @@ export default class Runtime {
     this._internalModuleRegistry = new Map();
     this._isCurrentlyExecutingManualMock = null;
     this._mainModule = null;
+    this._isolatedRequireInProgress = false;
     this._mockFactories = new Map();
     this._mockRegistry = new Map();
     this._moduleMockRegistry = new Map();
@@ -851,49 +853,58 @@ export default class Runtime {
       throw error;
     }
 
-    let moduleRegistry;
+    let disableIsolatedRequireInProgress = false;
+    let moduleRegistry = this._moduleRegistry;
 
     if (isInternal) {
       moduleRegistry = this._internalModuleRegistry;
-    } else if (this._isolatedModuleRegistry) {
+    } 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;
-    } else {
-      moduleRegistry = this._moduleRegistry;
     }
 
-    const module = moduleRegistry.get(modulePath);
-    if (module) {
-      return module.exports;
-    }
+    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,
+        path: path.dirname(modulePath),
+      };
+      moduleRegistry.set(modulePath, localModule);
 
-    // 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,
-      path: path.dirname(modulePath),
-    };
-    moduleRegistry.set(modulePath, localModule);
+      try {
+        this._loadModule(
+          localModule,
+          from,
+          moduleName,
+          modulePath,
+          options,
+          moduleRegistry,
+        );
+      } catch (error) {
+        moduleRegistry.delete(modulePath);
+        throw error;
+      }
 
-    try {
-      this._loadModule(
-        localModule,
-        from,
-        moduleName,
-        modulePath,
-        options,
-        moduleRegistry,
-      );
-    } catch (error) {
-      moduleRegistry.delete(modulePath);
-      throw error;
+      module = localModule;
     }
 
-    return localModule.exports;
+    if (disableIsolatedRequireInProgress) {
+      this._isolatedRequireInProgress = false;
+    }
+    return module.exports;
   }
 
   requireInternalModule<T = unknown>(from: string, to?: string): T {
@@ -931,18 +942,31 @@ export default class Runtime {
       {conditions: this.cjsConditions},
     );
 
-    if (this._isolatedMockRegistry?.has(moduleID)) {
-      return this._isolatedMockRegistry.get(moduleID);
-    } else if (this._mockRegistry.has(moduleID)) {
+    const isolatedMockRegistry =
+      !this._isolatedRequireInProgress && this._isolatedMockRegistry;
+    let disableIsolatedRequireInProgress = false;
+
+    if (isolatedMockRegistry) {
+      if (isolatedMockRegistry.get(moduleID)) {
+        return isolatedMockRegistry.get(moduleID);
+      } else {
+        this._isolatedRequireInProgress = true;
+        disableIsolatedRequireInProgress = true;
+      }
+    } else if (this._mockRegistry.get(moduleID)) {
       return this._mockRegistry.get(moduleID);
     }
 
-    const mockRegistry = this._isolatedMockRegistry || this._mockRegistry;
+    const mockRegistry = isolatedMockRegistry || this._mockRegistry;
 
     if (this._mockFactories.has(moduleID)) {
       // has check above makes this ok
       const module = this._mockFactories.get(moduleID)!();
       mockRegistry.set(moduleID, module);
+
+      if (disableIsolatedRequireInProgress) {
+        this._isolatedRequireInProgress = false;
+      }
       return module as T;
     }
 
@@ -1005,6 +1029,10 @@ export default class Runtime {
       mockRegistry.set(moduleID, this._generateMock(from, moduleName));
     }
 
+    if (disableIsolatedRequireInProgress) {
+      this._isolatedRequireInProgress = false;
+    }
+
     return mockRegistry.get(moduleID);
   }