fix(hoist-plugin): hoist pure constants to support experimental JSX transform in mocks (#10723)

Co-authored-by: Nicolò Ribaudo <nicolo.ribaudo@gmail.com>

Author: SimenB

CHANGELOG.md

@@ -7,6 +7,7 @@
 ### Fixes
 
 - `[babel-plugin-jest-hoist]` Preserve order of hoisted mock nodes within containing block ([#10536](https://github.com/facebook/jest/pull/10536))
+- `[babel-plugin-jest-hoist]` Hoist pure constants to support experimental JSX transform in hoisted mocks ([#10723](https://github.com/facebook/jest/pull/10723))
 - `[babel-preset-jest]` Update `babel-preset-current-node-syntax` to support top level await ([#10747](https://github.com/facebook/jest/pull/10747))
 - `[expect]` Stop modifying the sample in `expect.objectContaining()` ([#10711](https://github.com/facebook/jest/pull/10711))
 - `[jest-circus, jest-jasmine2]` fix: don't assume `stack` is always a string ([#10697](https://github.com/facebook/jest/pull/10697))

packages/babel-plugin-jest-hoist/package.json

@@ -20,9 +20,12 @@
   },
   "devDependencies": {
     "@babel/core": "^7.11.6",
+    "@babel/preset-react": "^7.12.1",
     "@types/babel__template": "^7.0.2",
     "@types/node": "*",
-    "babel-plugin-tester": "^10.0.0"
+    "@types/prettier": "^2.0.0",
+    "babel-plugin-tester": "^10.0.0",
+    "prettier": "^2.1.1"
   },
   "publishConfig": {
     "access": "public"

packages/babel-plugin-jest-hoist/src/__tests__/__snapshots__/hoistPlugin.test.ts.snap

@@ -1,5 +1,43 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
+exports[`babel-plugin-jest-hoist automatic react runtime: automatic react runtime 1`] = `
+
+jest.mock('./App', () => () => <div>Hello world</div>);
+
+      ↓ ↓ ↓ ↓ ↓ ↓
+
+var _jsxFileName = "/root/project/src/file.js";
+
+_getJestObj().mock("./App", () => () =>
+  /*#__PURE__*/ _jsxDEV(
+    "div",
+    {
+      children: "Hello world"
+    },
+    void 0,
+    false,
+    {
+      fileName: _jsxFileName,
+      lineNumber: 1,
+      columnNumber: 32
+    },
+    this
+  )
+);
+
+import { jsxDEV as _jsxDEV } from "react/jsx-dev-runtime";
+
+function _getJestObj() {
+  const { jest } = require("@jest/globals");
+
+  _getJestObj = () => jest;
+
+  return jest;
+}
+
+
+`;
+
 exports[`babel-plugin-jest-hoist top level mocking: top level mocking 1`] = `
 
 require('x');

packages/babel-plugin-jest-hoist/src/__tests__/hoistPlugin.test.ts

@@ -6,13 +6,41 @@
  *
  */
 
+import * as path from 'path';
 import pluginTester from 'babel-plugin-tester';
+import {format as formatCode} from 'prettier';
 import babelPluginJestHoist from '..';
 
 pluginTester({
   plugin: babelPluginJestHoist,
   pluginName: 'babel-plugin-jest-hoist',
   tests: {
+    'automatic react runtime': {
+      babelOptions: {
+        babelrc: false,
+        configFile: false,
+        filename: path.resolve(__dirname, '../file.js'),
+        presets: [
+          [
+            require.resolve('@babel/preset-react'),
+            {development: true, runtime: 'automatic'},
+          ],
+        ],
+      },
+      code: `
+        jest.mock('./App', () => () => <div>Hello world</div>);
+      `,
+      formatResult(code) {
+        // replace the filename with something that will be the same across OSes and machine
+        const codeWithoutSystemPath = code.replace(
+          /var _jsxFileName = ".*";/,
+          'var _jsxFileName = "/root/project/src/file.js";',
+        );
+
+        return formatCode(codeWithoutSystemPath, {parser: 'babel'});
+      },
+      snapshot: true,
+    },
     'top level mocking': {
       code: `
         require('x');

packages/babel-plugin-jest-hoist/src/index.ts

@@ -16,15 +16,20 @@ import {
   Identifier,
   Node,
   Program,
+  VariableDeclaration,
+  VariableDeclarator,
   callExpression,
   emptyStatement,
   isIdentifier,
+  variableDeclaration,
 } from '@babel/types';
 
 const JEST_GLOBAL_NAME = 'jest';
 const JEST_GLOBALS_MODULE_NAME = '@jest/globals';
 const JEST_GLOBALS_MODULE_JEST_EXPORT_NAME = 'jest';
 
+const hoistedVariables = new WeakSet<VariableDeclarator>();
+
 // We allow `jest`, `expect`, `require`, all default Node.js globals and all
 // ES2015 built-ins to be used inside of a `jest.mock` factory.
 // We also allow variables prefixed with `mock` as an escape-hatch.
@@ -133,12 +138,26 @@ FUNCTIONS.mock = args => {
       }
 
       if (!found) {
-        const isAllowedIdentifier =
+        let isAllowedIdentifier =
           (scope.hasGlobal(name) && ALLOWED_IDENTIFIERS.has(name)) ||
           /^mock/i.test(name) ||
           // Allow istanbul's coverage variable to pass.
           /^(?:__)?cov/.test(name);
 
+        if (!isAllowedIdentifier) {
+          const binding = scope.bindings[name];
+
+          if (binding?.path.isVariableDeclarator()) {
+            const {node} = binding.path;
+            const initNode = node.init;
+
+            if (initNode && binding.constant && scope.isPure(initNode, true)) {
+              hoistedVariables.add(node);
+              isAllowedIdentifier = true;
+            }
+          }
+        }
+
         if (!isAllowedIdentifier) {
           throw id.buildCodeFrameError(
             'The module factory of `jest.mock()` is not allowed to ' +
@@ -273,7 +292,7 @@ export default (): PluginObj<{
   visitor: {
     ExpressionStatement(exprStmt) {
       const jestObjExpr = extractJestObjExprIfHoistable(
-        exprStmt.get<'expression'>('expression'),
+        exprStmt.get('expression'),
       );
       if (jestObjExpr) {
         jestObjExpr.replaceWith(
@@ -285,24 +304,25 @@ export default (): PluginObj<{
   // in `post` to make sure we come after an import transform and can unshift above the `require`s
   post({path: program}) {
     const self = this;
+
     visitBlock(program);
-    program.traverse({
-      BlockStatement: visitBlock,
-    });
+    program.traverse({BlockStatement: visitBlock});
 
     function visitBlock(block: NodePath<BlockStatement> | NodePath<Program>) {
       // use a temporary empty statement instead of the real first statement, which may itself be hoisted
-      const [firstNonHoistedStatementOfBlock] = block.unshiftContainer(
-        'body',
+      const [varsHoistPoint, callsHoistPoint] = block.unshiftContainer('body', [
         emptyStatement(),
-      );
+        emptyStatement(),
+      ]);
       block.traverse({
         CallExpression: visitCallExpr,
+        VariableDeclarator: visitVariableDeclarator,
         // do not traverse into nested blocks, or we'll hoist calls in there out to this block
         // @ts-expect-error blacklist is not known
         blacklist: ['BlockStatement'],
       });
-      firstNonHoistedStatementOfBlock.remove();
+      callsHoistPoint.remove();
+      varsHoistPoint.remove();
 
       function visitCallExpr(callExpr: NodePath<CallExpression>) {
         const {
@@ -315,15 +335,32 @@ export default (): PluginObj<{
           const mockStmt = callExpr.getStatementParent();
 
           if (mockStmt) {
-            const mockStmtNode = mockStmt.node;
             const mockStmtParent = mockStmt.parentPath;
             if (mockStmtParent.isBlock()) {
+              const mockStmtNode = mockStmt.node;
               mockStmt.remove();
-              firstNonHoistedStatementOfBlock.insertBefore(mockStmtNode);
+              callsHoistPoint.insertBefore(mockStmtNode);
             }
           }
         }
       }
+
+      function visitVariableDeclarator(varDecl: NodePath<VariableDeclarator>) {
+        if (hoistedVariables.has(varDecl.node)) {
+          // should be assert function, but it's not. So let's cast below
+          varDecl.parentPath.assertVariableDeclaration();
+
+          const {kind, declarations} = varDecl.parent as VariableDeclaration;
+          if (declarations.length === 1) {
+            varDecl.parentPath.remove();
+          } else {
+            varDecl.remove();
+          }
+          varsHoistPoint.insertBefore(
+            variableDeclaration(kind, [varDecl.node]),
+          );
+        }
+      }
     }
   },
 });

yarn.lock

@@ -1456,7 +1456,7 @@ __metadata:
   languageName: node
   linkType: hard
 
-"@babel/preset-react@npm:*, @babel/preset-react@npm:^7.0.0, @babel/preset-react@npm:^7.9.4":
+"@babel/preset-react@npm:*, @babel/preset-react@npm:^7.0.0, @babel/preset-react@npm:^7.12.1, @babel/preset-react@npm:^7.9.4":
   version: 7.12.1
   resolution: "@babel/preset-react@npm:7.12.1"
   dependencies:
@@ -4830,13 +4830,16 @@ __metadata:
   resolution: "babel-plugin-jest-hoist@workspace:packages/babel-plugin-jest-hoist"
   dependencies:
     "@babel/core": ^7.11.6
+    "@babel/preset-react": ^7.12.1
     "@babel/template": ^7.3.3
     "@babel/types": ^7.3.3
     "@types/babel__core": ^7.0.0
     "@types/babel__template": ^7.0.2
     "@types/babel__traverse": ^7.0.6
     "@types/node": "*"
+    "@types/prettier": ^2.0.0
     babel-plugin-tester: ^10.0.0
+    prettier: ^2.1.1
   languageName: unknown
   linkType: soft