feat(runtime): add minimal support for mocking ESM (#11818)

Author: SimenB

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ### Features
 
+- `[jest-runtime]` Add experimental, limited (and undocumented) support for mocking ECMAScript Modules ([#11818](https://github.com/facebook/jest/pull/11818))
+
 ### Fixes
 
 - `[jest-types]` Export the `PrettyFormatOptions` interface ([#11801](https://github.com/facebook/jest/pull/11801))

e2e/__tests__/__snapshots__/nativeEsm.test.ts.snap

@@ -10,7 +10,7 @@ Ran all test suites matching /native-esm.tla.test.js/i.
 
 exports[`on node ^12.16.0 || >=13.7.0 runs test with native ESM 1`] = `
 Test Suites: 1 passed, 1 total
-Tests:       19 passed, 19 total
+Tests:       20 passed, 20 total
 Snapshots:   0 total
 Time:        <<REPLACED>>
 Ran all test suites matching /native-esm.test.js/i.

e2e/native-esm/__tests__/native-esm.test.js

@@ -177,3 +177,14 @@ test('require of ESM should throw correct error', () => {
     }),
   );
 });
+
+test('can mock module', async () => {
+  jestObject.unstable_mockModule('../mockedModule.mjs', () => ({foo: 'bar'}), {
+    virtual: true,
+  });
+
+  const importedMock = await import('../mockedModule.mjs');
+
+  expect(Object.keys(importedMock)).toEqual(['foo']);
+  expect(importedMock.foo).toEqual('bar');
+});

packages/jest-environment/src/index.ts

@@ -137,6 +137,14 @@ export interface Jest {
     moduleFactory?: () => unknown,
     options?: {virtual?: boolean},
   ): Jest;
+  /**
+   * Mocks a module with the provided module factory when it is being imported.
+   */
+  unstable_mockModule<T = unknown>(
+    moduleName: string,
+    moduleFactory: () => Promise<T> | T,
+    options?: {virtual?: boolean},
+  ): Jest;
   /**
    * Returns the actual module instead of a mock, bypassing all checks on
    * whether the module should receive a mock implementation or not.

packages/jest-runtime/src/index.ts

@@ -180,6 +180,7 @@ export default class Runtime {
   private _currentlyExecutingModulePath: string;
   private readonly _environment: JestEnvironment;
   private readonly _explicitShouldMock: Map<string, boolean>;
+  private readonly _explicitShouldMockModule: Map<string, boolean>;
   private _fakeTimersImplementation:
     | LegacyFakeTimers<unknown>
     | ModernFakeTimers
@@ -194,6 +195,8 @@ export default class Runtime {
   >;
   private _mockRegistry: Map<string, any>;
   private _isolatedMockRegistry: Map<string, any> | null;
+  private _moduleMockRegistry: Map<string, VMModule>;
+  private readonly _moduleMockFactories: Map<string, () => unknown>;
   private readonly _moduleMocker: ModuleMocker;
   private _isolatedModuleRegistry: ModuleRegistry | null;
   private _moduleRegistry: ModuleRegistry;
@@ -217,6 +220,7 @@ export default class Runtime {
   private readonly _transitiveShouldMock: Map<string, boolean>;
   private _unmockList: RegExp | undefined;
   private readonly _virtualMocks: Map<string, boolean>;
+  private readonly _virtualModuleMocks: Map<string, boolean>;
   private _moduleImplementation?: typeof nativeModule.Module;
   private readonly jestObjectCaches: Map<string, Jest>;
   private jestGlobals?: JestGlobals;
@@ -236,11 +240,14 @@ export default class Runtime {
     this._currentlyExecutingModulePath = '';
     this._environment = environment;
     this._explicitShouldMock = new Map();
+    this._explicitShouldMockModule = new Map();
     this._internalModuleRegistry = new Map();
     this._isCurrentlyExecutingManualMock = null;
     this._mainModule = null;
     this._mockFactories = new Map();
     this._mockRegistry = new Map();
+    this._moduleMockRegistry = new Map();
+    this._moduleMockFactories = new Map();
     invariant(
       this._environment.moduleMocker,
       '`moduleMocker` must be set on an environment when created',
@@ -260,6 +267,7 @@ export default class Runtime {
     this._fileTransforms = new Map();
     this._fileTransformsMutex = new Map();
     this._virtualMocks = new Map();
+    this._virtualModuleMocks = new Map();
     this.jestObjectCaches = new Map();
 
     this._mockMetaDataCache = new Map();
@@ -523,6 +531,16 @@ export default class Runtime {
 
     const [path, query] = specifier.split('?');
 
+    if (
+      this._shouldMock(
+        referencingIdentifier,
+        path,
+        this._explicitShouldMockModule,
+      )
+    ) {
+      return this.importMock(referencingIdentifier, path, context);
+    }
+
     const resolved = this._resolveModule(referencingIdentifier, path);
 
     if (
@@ -612,6 +630,46 @@ export default class Runtime {
     return evaluateSyntheticModule(module);
   }
 
+  private async importMock<T = unknown>(
+    from: Config.Path,
+    moduleName: string,
+    context: VMContext,
+  ): Promise<T> {
+    const moduleID = this._resolver.getModuleID(
+      this._virtualModuleMocks,
+      from,
+      moduleName,
+    );
+
+    if (this._moduleMockRegistry.has(moduleID)) {
+      return this._moduleMockRegistry.get(moduleID);
+    }
+
+    if (this._moduleMockFactories.has(moduleID)) {
+      const invokedFactory: any = await this._moduleMockFactories.get(
+        moduleID,
+        // has check above makes this ok
+      )!();
+
+      const module = new SyntheticModule(
+        Object.keys(invokedFactory),
+        function () {
+          Object.entries(invokedFactory).forEach(([key, value]) => {
+            // @ts-expect-error: TS doesn't know what `this` is
+            this.setExport(key, value);
+          });
+        },
+        {context, identifier: moduleName},
+      );
+
+      this._moduleMockRegistry.set(moduleID, module);
+
+      return evaluateSyntheticModule(module);
+    }
+
+    throw new Error('Attempting to import a mock without a factory');
+  }
+
   private getExportsOfCjs(modulePath: Config.Path) {
     const cachedNamedExports = this._cjsNamedExports.get(modulePath);
 
@@ -643,7 +701,7 @@ export default class Runtime {
     from: Config.Path,
     moduleName?: string,
     options?: InternalModuleOptions,
-    isRequireActual?: boolean | null,
+    isRequireActual = false,
   ): T {
     const moduleID = this._resolver.getModuleID(
       this._virtualMocks,
@@ -770,17 +828,12 @@ export default 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);
-    }
-
     const mockRegistry = this._isolatedMockRegistry || this._mockRegistry;
 
+    if (mockRegistry.get(moduleID)) {
+      return mockRegistry.get(moduleID);
+    }
+
     if (this._mockFactories.has(moduleID)) {
       // has check above makes this ok
       const module = this._mockFactories.get(moduleID)!();
@@ -896,7 +949,7 @@ export default class Runtime {
     }
 
     try {
-      if (this._shouldMock(from, moduleName)) {
+      if (this._shouldMock(from, moduleName, this._explicitShouldMock)) {
         return this.requireMock<T>(from, moduleName);
       } else {
         return this.requireModule<T>(from, moduleName);
@@ -952,6 +1005,7 @@ export default class Runtime {
     this._moduleRegistry.clear();
     this._esmoduleRegistry.clear();
     this._cjsNamedExports.clear();
+    this._moduleMockRegistry.clear();
 
     if (this._environment) {
       if (this._environment.global) {
@@ -1043,6 +1097,26 @@ export default class Runtime {
     this._mockFactories.set(moduleID, mockFactory);
   }
 
+  private setModuleMock(
+    from: string,
+    moduleName: string,
+    mockFactory: () => Promise<unknown> | unknown,
+    options?: {virtual?: boolean},
+  ): void {
+    if (options?.virtual) {
+      const mockPath = this._resolver.getModulePath(from, moduleName);
+
+      this._virtualModuleMocks.set(mockPath, true);
+    }
+    const moduleID = this._resolver.getModuleID(
+      this._virtualModuleMocks,
+      from,
+      moduleName,
+    );
+    this._explicitShouldMockModule.set(moduleID, true);
+    this._moduleMockFactories.set(moduleID, mockFactory);
+  }
+
   restoreAllMocks(): void {
     this._moduleMocker.restoreAllMocks();
   }
@@ -1063,12 +1137,15 @@ export default class Runtime {
     this._internalModuleRegistry.clear();
     this._mainModule = null;
     this._mockFactories.clear();
+    this._moduleMockFactories.clear();
     this._mockMetaDataCache.clear();
     this._shouldMockModuleCache.clear();
     this._shouldUnmockTransitiveDependenciesCache.clear();
     this._explicitShouldMock.clear();
+    this._explicitShouldMockModule.clear();
     this._transitiveShouldMock.clear();
     this._virtualMocks.clear();
+    this._virtualModuleMocks.clear();
     this._cacheFS.clear();
     this._unmockList = undefined;
 
@@ -1516,8 +1593,11 @@ export default class Runtime {
     );
   }
 
-  private _shouldMock(from: Config.Path, moduleName: string): boolean {
-    const explicitShouldMock = this._explicitShouldMock;
+  private _shouldMock(
+    from: Config.Path,
+    moduleName: string,
+    explicitShouldMock: Map<string, boolean>,
+  ): boolean {
     const moduleID = this._resolver.getModuleID(
       this._virtualMocks,
       from,
@@ -1687,6 +1767,18 @@ export default class Runtime {
       this.setMock(from, moduleName, mockFactory, options);
       return jestObject;
     };
+    const mockModule: Jest['unstable_mockModule'] = (
+      moduleName,
+      mockFactory,
+      options,
+    ) => {
+      if (typeof mockFactory !== 'function') {
+        throw new Error('`unstable_mockModule` must be passed a mock factory');
+      }
+
+      this.setModuleMock(from, moduleName, mockFactory, options);
+      return jestObject;
+    };
     const clearAllMocks = () => {
       this.clearAllMocks();
       return jestObject;
@@ -1821,6 +1913,7 @@ export default class Runtime {
       setTimeout,
       spyOn,
       unmock,
+      unstable_mockModule: mockModule,
       useFakeTimers,
       useRealTimers,
     };