Imported jest in mock (#13188)

jestjs · Aug 30, 2022 · 835a936 · 835a936
1 parent 2b04388
commit 835a936
Show file tree

Hide file tree

Showing 4 changed files with 187 additions and 2 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,6 +6,8 @@
 
 ### Fixes
 
+- `[babel-plugin-jest-hoist]` Support imported `jest` in mock factory ([#13188](https://github.com/facebook/jest/pull/13188))
+
 ### Chore & Maintenance
 
 ### Performance

diff --git a/packages/babel-plugin-jest-hoist/src/__tests__/__snapshots__/hoistPlugin.test.ts.snap b/packages/babel-plugin-jest-hoist/src/__tests__/__snapshots__/hoistPlugin.test.ts.snap
@@ -38,6 +38,113 @@ function _getJestObj() {
 }
 
 
+`;
+
+exports[`babel-plugin-jest-hoist global jest.mock within jest.mock: global jest.mock within jest.mock 1`] = `
+
+jest.mock('some-module', () => {
+  jest.mock('some-module');
+});
+
+      ↓ ↓ ↓ ↓ ↓ ↓
+
+_getJestObj().mock('some-module', () => {
+  _getJestObj().mock('some-module');
+});
+
+function _getJestObj() {
+  const {jest} = require('@jest/globals');
+
+  _getJestObj = () => jest;
+
+  return jest;
+}
+
+
+`;
+
+exports[`babel-plugin-jest-hoist global jest.requireActual in jest.mock: global jest.requireActual in jest.mock 1`] = `
+
+jest.mock('some-module', () => {
+  jest.requireActual('some-module');
+});
+
+jest.requireActual('some-module');
+
+      ↓ ↓ ↓ ↓ ↓ ↓
+
+_getJestObj().mock('some-module', () => {
+  _getJestObj().requireActual('some-module');
+});
+
+function _getJestObj() {
+  const {jest} = require('@jest/globals');
+
+  _getJestObj = () => jest;
+
+  return jest;
+}
+
+jest.requireActual('some-module');
+
+
+`;
+
+exports[`babel-plugin-jest-hoist imported jest.mock within jest.mock: imported jest.mock within jest.mock 1`] = `
+
+import {jest} from '@jest/globals';
+
+jest.mock('some-module', () => {
+  jest.mock('some-module');
+});
+
+      ↓ ↓ ↓ ↓ ↓ ↓
+
+_getJestObj().mock('some-module', () => {
+  _getJestObj().mock('some-module');
+});
+
+function _getJestObj() {
+  const {jest} = require('@jest/globals');
+
+  _getJestObj = () => jest;
+
+  return jest;
+}
+
+import {jest} from '@jest/globals';
+
+
+`;
+
+exports[`babel-plugin-jest-hoist imported jest.requireActual in jest.mock: imported jest.requireActual in jest.mock 1`] = `
+
+import {jest} from '@jest/globals';
+
+jest.mock('some-module', () => {
+  jest.requireActual('some-module');
+});
+
+jest.requireActual('some-module');
+
+      ↓ ↓ ↓ ↓ ↓ ↓
+
+_getJestObj().mock('some-module', () => {
+  _getJestObj().requireActual('some-module');
+});
+
+function _getJestObj() {
+  const {jest} = require('@jest/globals');
+
+  _getJestObj = () => jest;
+
+  return jest;
+}
+
+import {jest} from '@jest/globals';
+jest.requireActual('some-module');
+
+
 `;
 
 exports[`babel-plugin-jest-hoist required \`jest\` within \`jest\`: required \`jest\` within \`jest\` 1`] = `

diff --git a/packages/babel-plugin-jest-hoist/src/__tests__/hoistPlugin.test.ts b/packages/babel-plugin-jest-hoist/src/__tests__/hoistPlugin.test.ts
@@ -90,6 +90,50 @@ pluginTester({
       formatResult,
       snapshot: true,
     },
+    'imported jest.mock within jest.mock': {
+      code: formatResult(`
+        import {jest} from '@jest/globals';
+
+        jest.mock('some-module', () => {
+          jest.mock('some-module');
+        });
+      `),
+      formatResult,
+      snapshot: true,
+    },
+    'global jest.mock within jest.mock': {
+      code: formatResult(`
+        jest.mock('some-module', () => {
+          jest.mock('some-module');
+        });
+      `),
+      formatResult,
+      snapshot: true,
+    },
+    'imported jest.requireActual in jest.mock': {
+      code: formatResult(`
+        import {jest} from '@jest/globals';
+
+        jest.mock('some-module', () => {
+          jest.requireActual('some-module');
+        });
+
+        jest.requireActual('some-module');
+      `),
+      formatResult,
+      snapshot: true,
+    },
+    'global jest.requireActual in jest.mock': {
+      code: formatResult(`
+        jest.mock('some-module', () => {
+          jest.requireActual('some-module');
+        });
+
+        jest.requireActual('some-module');
+      `),
+      formatResult,
+      snapshot: true,
+    },
   },
   /* eslint-enable */
 });
diff --git a/packages/babel-plugin-jest-hoist/src/index.ts b/packages/babel-plugin-jest-hoist/src/index.ts
@@ -14,6 +14,7 @@ import {
   CallExpression,
   Expression,
   Identifier,
+  ImportDeclaration,
   MemberExpression,
   Node,
   Program,
@@ -31,6 +32,7 @@ const JEST_GLOBALS_MODULE_NAME = '@jest/globals';
 const JEST_GLOBALS_MODULE_JEST_EXPORT_NAME = 'jest';
 
 const hoistedVariables = new WeakSet<VariableDeclarator>();
+const hoistedJestExpressions = new WeakSet<Expression>();
 
 // We allow `jest`, `expect`, `require`, all default Node.js globals and all
 // ES2015 built-ins to be used inside of a `jest.mock` factory.
@@ -161,6 +163,19 @@ FUNCTIONS.mock = args => {
               hoistedVariables.add(node);
               isAllowedIdentifier = true;
             }
+          } else if (binding?.path.isImportSpecifier()) {
+            const importDecl = binding.path
+              .parentPath as NodePath<ImportDeclaration>;
+            const imported = binding.path.node.imported;
+            if (
+              importDecl.node.source.value === JEST_GLOBALS_MODULE_NAME &&
+              (isIdentifier(imported) ? imported.name : imported.value) ===
+                JEST_GLOBALS_MODULE_JEST_EXPORT_NAME
+            ) {
+              isAllowedIdentifier = true;
+              // Imports are already hoisted, so we don't need to add it
+              // to hoistedVariables.
+            }
           }
         }
 
@@ -264,9 +279,26 @@ const extractJestObjExprIfHoistable = (
   // Important: Call the function check last
   // It might throw an error to display to the user,
   // which should only happen if we're already sure it's a call on the Jest object.
-  const functionLooksHoistable = FUNCTIONS[propertyName]?.(args);
+  let functionLooksHoistableOrInHoistable = FUNCTIONS[propertyName]?.(args);
+
+  for (
+    let path: NodePath<Node> | null = expr;
+    path && !functionLooksHoistableOrInHoistable;
+    path = path.parentPath
+  ) {
+    functionLooksHoistableOrInHoistable = hoistedJestExpressions.has(
+      // @ts-expect-error: it's ok if path.node is not an Expression, .has will
+      // just return false.
+      path.node,
+    );
+  }
+
+  if (functionLooksHoistableOrInHoistable) {
+    hoistedJestExpressions.add(expr.node);
+    return jestObjExpr;
+  }
 
-  return functionLooksHoistable ? jestObjExpr : null;
+  return null;
 };
 
 /* eslint-disable sort-keys */