fix: hoist imports of @jest/globals correctly

e2e/__tests__/babelPluginJestHoist.test.ts

@@ -15,8 +15,8 @@ beforeEach(() => {
   run('yarn', DIR);
 });
 
-it('sucessfully runs the tests inside `babel-plugin-jest-hoist/`', () => {
+it('successfully runs the tests inside `babel-plugin-jest-hoist/`', () => {
   const {json} = runWithJson(DIR, ['--no-cache', '--coverage']);
   expect(json.success).toBe(true);
-  expect(json.numTotalTestSuites).toBe(3);
+  expect(json.numTotalTestSuites).toBe(4);
 });

e2e/babel-plugin-jest-hoist/__tests__/importJest.test.js

@@ -0,0 +1,53 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ */
+
+/* eslint-disable import/no-duplicates */
+import {jest} from '@jest/globals';
+import {jest as aliasedJest} from '@jest/globals';
+import * as JestGlobals from '@jest/globals';
+/* eslint-enable import/no-duplicates */
+
+import a from '../__test_modules__/a';
+import b from '../__test_modules__/b';
+import c from '../__test_modules__/c';
+import d from '../__test_modules__/d';
+
+// These will be hoisted above imports
+
+jest.unmock('../__test_modules__/a');
+aliasedJest.unmock('../__test_modules__/b');
+JestGlobals.jest.unmock('../__test_modules__/c');
+
+// These will not be hoisted above imports
+
+{
+  const jest = {unmock: () => {}};
+  jest.unmock('../__test_modules__/d');
+}
+
+// tests
+
+test('named import', () => {
+  expect(a._isMockFunction).toBe(undefined);
+  expect(a()).toBe('unmocked');
+});
+
+test('aliased named import', () => {
+  expect(b._isMockFunction).toBe(undefined);
+  expect(b()).toBe('unmocked');
+});
+
+test('namespace import', () => {
+  expect(c._isMockFunction).toBe(undefined);
+  expect(c()).toBe('unmocked');
+});
+
+test('fake jest, shadowed import', () => {
+  expect(d._isMockFunction).toBe(true);
+  expect(d()).toBe(undefined);
+});

e2e/babel-plugin-jest-hoist/__tests__/integration.test.js

@@ -15,7 +15,7 @@ import a from '../__test_modules__/a';
 import b from '../__test_modules__/b';
 import c from '../__test_modules__/c';
 import d from '../__test_modules__/d';
-import e from '../__test_modules__/e';
+import f from '../__test_modules__/f';
 import jestBackticks from '../__test_modules__/jestBackticks';
 
 // The virtual mock call below will be hoisted above this `require` call.
@@ -25,7 +25,16 @@ const virtualModule = require('virtual-module');
 jest.unmock('react');
 jest.deepUnmock('../__test_modules__/Unmocked');
 jest.unmock('../__test_modules__/c').unmock('../__test_modules__/d');
-jest.mock('../__test_modules__/e', () => {
+
+let e;
+(function () {
+  const _getJestObj = 42;
+  e = require('../__test_modules__/e').default;
+  // hoisted to the top of the function scope
+  jest.unmock('../__test_modules__/e');
+})();
+
+jest.mock('../__test_modules__/f', () => {
   if (!global.CALLS) {
     global.CALLS = 0;
   }
@@ -52,8 +61,13 @@ jest.mock('has-flow-types', () => (props: {children: mixed}) => 3, {
 // These will not be hoisted
 jest.unmock('../__test_modules__/a').dontMock('../__test_modules__/b');
 // eslint-disable-next-line no-useless-concat
-jest.unmock('../__test_modules__/' + 'c');
+jest.unmock('../__test_modules__/' + 'a');
 jest.dontMock('../__test_modules__/Mocked');
+{
+  const jest = {unmock: () => {}};
+  // Would error (used before initialization) if hoisted to the top of the scope
+  jest.unmock('../__test_modules__/a');
+}
 
 // This must not throw an error
 const myObject = {mock: () => {}};
@@ -84,14 +98,17 @@ describe('babel-plugin-jest-hoist', () => {
 
     expect(d._isMockFunction).toBe(undefined);
     expect(d()).toEqual('unmocked');
+
+    expect(e._isMock).toBe(undefined);
+    expect(e()).toEqual('unmocked');
   });
 
   it('hoists mock call with 2 arguments', () => {
     const path = require('path');
 
-    expect(e._isMock).toBe(true);
+    expect(f._isMock).toBe(true);
 
-    const mockFn = e.fn();
+    const mockFn = f.fn();
     expect(mockFn()).toEqual([path.sep, undefined, undefined]);
   });
 
@@ -100,10 +117,10 @@ describe('babel-plugin-jest-hoist', () => {
 
     global.CALLS = 0;
 
-    require('../__test_modules__/e');
+    require('../__test_modules__/f');
     expect(global.CALLS).toEqual(1);
 
-    require('../__test_modules__/e');
+    require('../__test_modules__/f');
     expect(global.CALLS).toEqual(1);
 
     delete global.CALLS;

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

@@ -20,10 +20,11 @@
     }
   },
   "dependencies": {
+    "@babel/template": "^7.3.3",
+    "@babel/types": "^7.3.3",
     "@types/babel__traverse": "^7.0.6"
   },
   "devDependencies": {
-    "@babel/types": "^7.3.3",
     "@types/node": "*"
   },
   "publishConfig": {

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

@@ -6,8 +6,21 @@
  *
  */
 
-import type {NodePath, Visitor} from '@babel/traverse';
-import type {Identifier} from '@babel/types';
+import type {NodePath} from '@babel/traverse';
+import {
+  Expression,
+  Identifier,
+  Node,
+  Program,
+  callExpression,
+  isIdentifier,
+} from '@babel/types';
+import {statement} from '@babel/template';
+import type {PluginObj} from '@babel/core';
+
+const JEST_GLOBAL_NAME = 'jest';
+const JEST_GLOBALS_MODULE_NAME = '@jest/globals';
+const JEST_GLOBALS_MODULE_JEST_EXPORT_NAME = 'jest';
 
 // We allow `jest`, `expect`, `require`, all default Node.js globals and all
 // ES2015 built-ins to be used inside of a `jest.mock` factory.
@@ -70,19 +83,19 @@ const WHITELISTED_IDENTIFIERS = new Set<string>(
   ].sort(),
 );
 
-const JEST_GLOBAL = {name: 'jest'};
-// TODO: Should be Visitor<{ids: Set<NodePath<Identifier>>}>, but `ReferencedIdentifier` doesn't exist
 const IDVisitor = {
-  ReferencedIdentifier(path: NodePath<Identifier>) {
-    // @ts-ignore: passed as Visitor State
-    this.ids.add(path);
+  ReferencedIdentifier(
+    path: NodePath<Identifier>,
+    {ids}: {ids: Set<NodePath<Identifier>>},
+  ) {
+    ids.add(path);
   },
   blacklist: ['TypeAnnotation', 'TSTypeAnnotation', 'TSTypeReference'],
 };
 
 const FUNCTIONS: Record<
   string,
-  (args: Array<NodePath>) => boolean
+  <T extends Node>(args: Array<NodePath<T>>) => boolean
 > = Object.create(null);
 
 FUNCTIONS.mock = args => {
@@ -100,7 +113,7 @@ FUNCTIONS.mock = args => {
 
     const ids: Set<NodePath<Identifier>> = new Set();
     const parentScope = moduleFactory.parentPath.scope;
-    // @ts-ignore: Same as above: ReferencedIdentifier doesn't exist
+    // @ts-ignore: ReferencedIdentifier is not known on visitors
     moduleFactory.traverse(IDVisitor, {ids});
     for (const id of ids) {
       const {name} = id.node;
@@ -152,38 +165,140 @@ FUNCTIONS.deepUnmock = args => args.length === 1 && args[0].isStringLiteral();
 FUNCTIONS.disableAutomock = FUNCTIONS.enableAutomock = args =>
   args.length === 0;
 
-export default (): {visitor: Visitor} => {
-  const shouldHoistExpression = (expr: NodePath): boolean => {
-    if (!expr.isCallExpression()) {
-      return false;
-    }
+const createJestObjectGetter = statement`
+function GETTER_NAME() {
+  const { JEST_GLOBALS_MODULE_JEST_EXPORT_NAME } = require("JEST_GLOBALS_MODULE_NAME");
+  GETTER_NAME = () => JEST_GLOBALS_MODULE_JEST_EXPORT_NAME;
+  return JEST_GLOBALS_MODULE_JEST_EXPORT_NAME;
+}
+`;
 
-    // TODO: avoid type casts - the types can be arrays (is it possible to ignore that without casting?)
-    const callee = expr.get('callee') as NodePath;
-    const expressionArguments = expr.get('arguments');
-    const object = callee.get('object') as NodePath;
-    const property = callee.get('property') as NodePath;
-    return (
-      property.isIdentifier() &&
-      FUNCTIONS[property.node.name] &&
-      (object.isIdentifier(JEST_GLOBAL) ||
-        (callee.isMemberExpression() && shouldHoistExpression(object))) &&
-      FUNCTIONS[property.node.name](
-        Array.isArray(expressionArguments)
-          ? expressionArguments
-          : [expressionArguments],
-      )
-    );
-  };
-
-  const visitor: Visitor = {
-    ExpressionStatement(path) {
-      if (shouldHoistExpression(path.get('expression') as NodePath)) {
-        // @ts-ignore: private, magical property
-        path.node._blockHoist = Infinity;
-      }
-    },
-  };
+const isJestObject = (expression: NodePath<Expression>): boolean => {
+  // global
+  if (
+    expression.isIdentifier() &&
+    expression.node.name === JEST_GLOBAL_NAME &&
+    !expression.scope.hasBinding(JEST_GLOBAL_NAME)
+  ) {
+    return true;
+  }
+  // import { jest } from '@jest/globals'
+  if (
+    expression.referencesImport(
+      JEST_GLOBALS_MODULE_NAME,
+      JEST_GLOBALS_MODULE_JEST_EXPORT_NAME,
+    )
+  ) {
+    return true;
+  }
+  // import * as JestGlobals from '@jest/globals'
+  if (
+    expression.isMemberExpression() &&
+    !expression.node.computed &&
+    expression
+      .get<'object'>('object')
+      .referencesImport(JEST_GLOBALS_MODULE_NAME, '*') &&
+    expression.node.property.name === JEST_GLOBALS_MODULE_JEST_EXPORT_NAME
+  ) {
+    return true;
+  }
 
-  return {visitor};
+  return false;
 };
+
+const extractJestObjExprIfHoistable = <T extends Node>(
+  expr: NodePath<T>,
+): NodePath<Expression> | null => {
+  if (!expr.isCallExpression()) {
+    return null;
+  }
+
+  const callee = expr.get<'callee'>('callee');
+  const args = expr.get<'arguments'>('arguments');
+
+  if (!callee.isMemberExpression() || callee.node.computed) {
+    return null;
+  }
+
+  const object = callee.get<'object'>('object');
+  const property = callee.get<'property'>('property') as NodePath<Identifier>;
+  const propertyName = property.node.name;
+
+  const jestObjExpr = isJestObject(object)
+    ? object
+    : // The Jest object could be returned from another call since the functions are all chainable.
+      extractJestObjExprIfHoistable(object);
+  if (!jestObjExpr) {
+    return null;
+  }
+
+  // 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);
+
+  return functionLooksHoistable ? jestObjExpr : null;
+};
+
+/* eslint-disable sort-keys,@typescript-eslint/explicit-module-boundary-types */
+export default (): PluginObj<{
+  declareJestObjGetterIdentifier: () => Identifier;
+  jestObjGetterIdentifier?: Identifier;
+}> => ({
+  pre({path: program}: {path: NodePath<Program>}) {
+    this.declareJestObjGetterIdentifier = () => {
+      if (this.jestObjGetterIdentifier) {
+        return this.jestObjGetterIdentifier;
+      }
+
+      this.jestObjGetterIdentifier = program.scope.generateUidIdentifier(
+        'getJestObj',
+      );
+
+      program.unshiftContainer('body', [
+        createJestObjectGetter({
+          GETTER_NAME: this.jestObjGetterIdentifier.name,
+          JEST_GLOBALS_MODULE_JEST_EXPORT_NAME,
+          JEST_GLOBALS_MODULE_NAME,
+        }),
+      ]);
+
+      return this.jestObjGetterIdentifier;
+    };
+  },
+  visitor: {
+    ExpressionStatement(exprStmt) {
+      const jestObjExpr = extractJestObjExprIfHoistable(
+        exprStmt.get<'expression'>('expression'),
+      );
+      if (jestObjExpr) {
+        jestObjExpr.replaceWith(
+          callExpression(this.declareJestObjGetterIdentifier(), []),
+        );
+      }
+    },
+  },
+  // in `post` to make sure we come after an import transform and can unshift above the `require`s
+  post({path: program}: {path: NodePath<Program>}) {
+    program.traverse({
+      CallExpression: callExpr => {
+        const {
+          node: {callee},
+        } = callExpr;
+        if (
+          isIdentifier(callee) &&
+          callee.name === this.jestObjGetterIdentifier?.name
+        ) {
+          const mockStmt = callExpr.getStatementParent();
+          const mockStmtNode = mockStmt.node;
+          const mockStmtParent = mockStmt.parentPath;
+          if (mockStmtParent.isBlock()) {
+            mockStmt.remove();
+            mockStmtParent.unshiftContainer('body', [mockStmtNode]);
+          }
+        }
+      },
+    });
+  },
+});
+/* eslint-enable sort-keys,@typescript-eslint/explicit-module-boundary-types */

yarn.lock

@@ -994,7 +994,7 @@
   dependencies:
     regenerator-runtime "^0.13.4"
 
-"@babel/template@^7.0.0", "@babel/template@^7.7.4", "@babel/template@^7.8.3", "@babel/template@^7.8.6":
+"@babel/template@^7.0.0", "@babel/template@^7.3.3", "@babel/template@^7.7.4", "@babel/template@^7.8.3", "@babel/template@^7.8.6":
   version "7.8.6"
   resolved "https://registry.yarnpkg.com/@babel/template/-/template-7.8.6.tgz#86b22af15f828dfb086474f964dcc3e39c43ce2b"
   integrity sha512-zbMsPMy/v0PWFZEhQJ66bqjhH+z0JgMoBWuikXybgG3Gkd/3t5oQ1Rw2WQhnSrsOmsKXnZOx15tkC4qON/+JPg==