feat(runtime): add jest.mockModule

Author: SimenB

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:       18 passed, 18 total
+Tests:       19 passed, 19 total
 Snapshots:   0 total
 Time:        <<REPLACED>>
 Ran all test suites matching /native-esm.test.js/i.

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

@@ -166,3 +166,14 @@ test('handle circular dependency', async () => {
   expect(moduleA.moduleB.id).toBe('circularDependentB');
   expect(moduleA.moduleB.moduleA).toBe(moduleA);
 });
+
+test('can mock module', async () => {
+  jestObject.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

@@ -146,6 +146,14 @@ export interface Jest {
     moduleFactory?: () => unknown,
     options?: {virtual?: boolean},
   ): Jest;
+  /**
+   * Mocks a module with an auto-mocked version when it is being required.
+   */
+  mockModule(
+    moduleName: string,
+    moduleFactory?: () => Promise<unknown> | unknown,
+    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-resolve/src/index.ts

@@ -313,10 +313,8 @@ class Resolver {
   getModuleID(
     virtualMocks: Map<string, boolean>,
     from: Config.Path,
-    _moduleName?: string,
+    moduleName = '',
   ): string {
-    const moduleName = _moduleName || '';
-
     const key = from + path.delimiter + moduleName;
     const cachedModuleID = this._moduleIDCache.get(key);
     if (cachedModuleID) {

packages/jest-runtime/src/index.ts

@@ -148,6 +148,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
@@ -162,6 +163,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;
@@ -183,6 +186,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;
@@ -201,11 +205,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',
@@ -223,6 +230,7 @@ export default class Runtime {
     this._sourceMapRegistry = new Map();
     this._fileTransforms = new Map();
     this._virtualMocks = new Map();
+    this._virtualModuleMocks = new Map();
     this.jestObjectCaches = new Map();
 
     this._mockMetaDataCache = new Map();
@@ -490,6 +498,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 (
@@ -505,6 +523,8 @@ export default class Runtime {
   async unstable_importModule(
     from: Config.Path,
     moduleName?: string,
+    // TODO: implement this
+    _isImportActual = false,
   ): Promise<void> {
     invariant(
       runtimeSupportsVmModules,
@@ -552,6 +572,109 @@ 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);
+          });
+        },
+        // should identifier be `node://${moduleName}`?
+        {context, identifier: moduleName},
+      );
+
+      this._moduleMockRegistry.set(moduleID, module);
+
+      return evaluateSyntheticModule(module);
+    }
+
+    const manualMockOrStub = this._resolver.getMockModule(from, moduleName);
+
+    let modulePath =
+      this._resolver.getMockModule(from, moduleName) ||
+      this._resolveModule(from, moduleName);
+
+    let isManualMock =
+      manualMockOrStub &&
+      !this._resolver.resolveStubModuleName(from, moduleName);
+    if (!isManualMock) {
+      // If the actual module file has a __mocks__ dir sitting immediately next
+      // to it, look to see if there is a manual mock for this file.
+      //
+      // subDir1/my_module.js
+      // subDir1/__mocks__/my_module.js
+      // subDir2/my_module.js
+      // subDir2/__mocks__/my_module.js
+      //
+      // Where some other module does a relative require into each of the
+      // respective subDir{1,2} directories and expects a manual mock
+      // corresponding to that particular my_module.js file.
+
+      const moduleDir = path.dirname(modulePath);
+      const moduleFileName = path.basename(modulePath);
+      const potentialManualMock = path.join(
+        moduleDir,
+        '__mocks__',
+        moduleFileName,
+      );
+      if (fs.existsSync(potentialManualMock)) {
+        isManualMock = true;
+        modulePath = potentialManualMock;
+      }
+    }
+    if (isManualMock) {
+      const localModule: InitialModule = {
+        children: [],
+        exports: {},
+        filename: modulePath,
+        id: modulePath,
+        loaded: false,
+        path: modulePath,
+      };
+
+      this._loadModule(
+        localModule,
+        from,
+        moduleName,
+        modulePath,
+        undefined,
+        this._moduleMockRegistry,
+      );
+
+      this._moduleMockRegistry.set(moduleID, localModule.exports);
+    } else {
+      // Look for a real module to generate an automock from
+      this._moduleMockRegistry.set(
+        moduleID,
+        this._generateMock(from, moduleName),
+      );
+    }
+
+    return this._moduleMockRegistry.get(moduleID);
+  }
+
   private getExportsOfCjs(modulePath: Config.Path) {
     const cachedNamedExports = this._cjsNamedExports.get(modulePath);
 
@@ -583,7 +706,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,
@@ -620,12 +743,10 @@ export default class Runtime {
 
     if (options?.isInternalModule) {
       moduleRegistry = this._internalModuleRegistry;
+    } else if (this._isolatedModuleRegistry) {
+      moduleRegistry = this._isolatedModuleRegistry;
     } else {
-      if (this._isolatedModuleRegistry) {
-        moduleRegistry = this._isolatedModuleRegistry;
-      } else {
-        moduleRegistry = this._moduleRegistry;
-      }
+      moduleRegistry = this._moduleRegistry;
     }
 
     const module = moduleRegistry.get(modulePath);
@@ -686,17 +807,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)!();
@@ -813,7 +929,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);
@@ -869,6 +985,7 @@ export default class Runtime {
     this._moduleRegistry.clear();
     this._esmoduleRegistry.clear();
     this._cjsNamedExports.clear();
+    this._moduleMockRegistry.clear();
 
     if (this._environment) {
       if (this._environment.global) {
@@ -957,6 +1074,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();
   }
@@ -977,12 +1114,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;
 
@@ -1374,8 +1514,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,
@@ -1543,6 +1686,24 @@ export default class Runtime {
       this.setMock(from, moduleName, mockFactory, options);
       return jestObject;
     };
+    const mockModule: Jest['mockModule'] = (
+      moduleName,
+      mockFactory,
+      options,
+    ) => {
+      if (mockFactory !== undefined) {
+        this.setModuleMock(from, moduleName, mockFactory, options);
+        return jestObject;
+      }
+
+      const moduleID = this._resolver.getModuleID(
+        this._virtualMocks,
+        from,
+        moduleName,
+      );
+      this._explicitShouldMockModule.set(moduleID, true);
+      return jestObject;
+    };
     const clearAllMocks = () => {
       this.clearAllMocks();
       return jestObject;
@@ -1641,6 +1802,7 @@ export default class Runtime {
       isMockFunction: this._moduleMocker.isMockFunction,
       isolateModules,
       mock,
+      mockModule,
       requireActual: this.requireActual.bind(this, from),
       requireMock: this.requireMock.bind(this, from),
       resetAllMocks,