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:       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.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 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-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 (
@@ -563,6 +581,8 @@ export default class Runtime {
   async unstable_importModule(
     from: Config.Path,
     moduleName?: string,
+    // TODO: implement this
+    _isImportActual = false,
   ): Promise<void> {
     invariant(
       runtimeSupportsVmModules,
@@ -612,6 +632,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);
 
@@ -643,7 +766,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 +893,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 +1014,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 +1070,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 +1162,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 +1202,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 +1658,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 +1832,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;
@@ -1785,6 +1948,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,