alangpierce · alangpierce · Apr 11, 2021 · Apr 11, 2021
diff --git a/src/transformers/JestHoistTransformer.ts b/src/transformers/JestHoistTransformer.ts
@@ -1,4 +1,5 @@
 import type CJSImportProcessor from "../CJSImportProcessor";
+import type NameManager from "../NameManager";
 import {TokenType as tt} from "../parser/tokenizer/types";
 import type TokenProcessor from "../TokenProcessor";
 import type RootTransformer from "./RootTransformer";
@@ -10,13 +11,18 @@ const HOISTED_METHODS = ["mock", "unmock", "enableAutomock", "disableAutomock"];
 /**
  * Implementation of babel-plugin-jest-hoist, which hoists up some jest method
  * calls above the imports to allow them to override other imports.
+ *
+ * To preserve line numbers, rather than directly moving the jest.mock code, we
+ * wrap each invocation in a function statement and then call the function from
+ * the top of the file.
  */
 export default class JestHoistTransformer extends Transformer {
-  private readonly hoistedCalls: Array<string> = [];
+  private readonly hoistedFunctionNames: Array<string> = [];
 
   constructor(
     readonly rootTransformer: RootTransformer,
     readonly tokens: TokenProcessor,
+    readonly nameManager: NameManager,
     readonly importProcessor: CJSImportProcessor | null,
   ) {
     super();
@@ -40,10 +46,10 @@ export default class JestHoistTransformer extends Transformer {
   }
 
   getHoistedCode(): string {
-    if (this.hoistedCalls.length > 0) {
+    if (this.hoistedFunctionNames.length > 0) {
       // This will be placed before module interop code, but that's fine since
       // imports aren't allowed in module mock factories.
-      return `\n${JEST_GLOBAL_NAME}${this.hoistedCalls.join("")};`;
+      return this.hoistedFunctionNames.map((name) => `${name}();`).join("");
     }
     return "";
   }
@@ -57,46 +63,46 @@ export default class JestHoistTransformer extends Transformer {
    * We do not apply the same checks of the arguments as babel-plugin-jest-hoist does.
    */
   private extractHoistedCalls(): boolean {
-    // We remove the `jest` expression, then add it back later if we find a non-hoisted call
+    // We're handling a chain of calls where `jest` may or may not need to be inserted for each call
+    // in the chain, so remove the initial `jest` to make the loop implementation cleaner.
     this.tokens.removeToken();
-    let restoredJest = false;
+    // Track some state so that multiple non-hoisted chained calls in a row keep their chaining
+    // syntax.
+    let followsNonHoistedJestCall = false;
 
-    // Iterate through all chained calls on the jest object
+    // Iterate through all chained calls on the jest object.
     while (this.tokens.matches3(tt.dot, tt.name, tt.parenL)) {
       const methodName = this.tokens.identifierNameAtIndex(this.tokens.currentIndex() + 1);
       const shouldHoist = HOISTED_METHODS.includes(methodName);
       if (shouldHoist) {
-        // We've matched e.g. `.mock(...)` or similar call
-        // Start by applying transforms to the entire call, including parameters
-        const snapshotBefore = this.tokens.snapshot();
-        this.tokens.copyToken();
+        // We've matched e.g. `.mock(...)` or similar call.
+        // Replace the initial `.` with `function __jestHoist(){jest.`
+        const hoistedFunctionName = this.nameManager.claimFreeName("__jestHoist");
+        this.hoistedFunctionNames.push(hoistedFunctionName);
+        this.tokens.replaceToken(`function ${hoistedFunctionName}(){${JEST_GLOBAL_NAME}.`);
         this.tokens.copyToken();
         this.tokens.copyToken();
         this.rootTransformer.processBalancedCode();
         this.tokens.copyExpectedToken(tt.parenR);
-        const snapshotAfter = this.tokens.snapshot();
-
-        // Then grab the transformed code and store it for hoisting
-        const processedCall = snapshotAfter.resultCode.slice(snapshotBefore.resultCode.length);
-        this.hoistedCalls.push(processedCall);
-
-        // Now go back and remove the entire method call
-        const endIndex = this.tokens.currentIndex();
-        this.tokens.restoreToSnapshot(snapshotBefore);
-        while (this.tokens.currentIndex() < endIndex) {
-          this.tokens.removeToken();
-        }
+        this.tokens.appendCode(";}");
+        followsNonHoistedJestCall = false;
       } else {
-        if (!restoredJest) {
-          restoredJest = true;
-          this.tokens.appendCode(JEST_GLOBAL_NAME);
+        // This is a non-hoisted method, so just transform the code as usual.
+        if (followsNonHoistedJestCall) {
+          // If we didn't hoist the previous call, we can leave the code as-is to chain off of the
+          // previous method call. It's important to preserve the code here because we don't know
+          // for sure that the method actually returned the jest object for chaining.
+          this.tokens.copyToken();
+        } else {
+          // If we hoisted the previous call, we know it returns the jest object back, so we insert
+          // the identifier `jest` to continue the chain.
+          this.tokens.replaceToken(`${JEST_GLOBAL_NAME}.`);
         }
-        // When not hoisting we just transform the method call as usual
-        this.tokens.copyToken();
         this.tokens.copyToken();
         this.tokens.copyToken();
         this.rootTransformer.processBalancedCode();
         this.tokens.copyExpectedToken(tt.parenR);
+        followsNonHoistedJestCall = true;
       }
     }
 

diff --git a/src/transformers/RootTransformer.ts b/src/transformers/RootTransformer.ts
@@ -102,7 +102,9 @@ export default class RootTransformer {
       );
     }
     if (transforms.includes("jest")) {
-      this.transformers.push(new JestHoistTransformer(this, tokenProcessor, importProcessor));
+      this.transformers.push(
+        new JestHoistTransformer(this, tokenProcessor, this.nameManager, importProcessor),
+      );
     }
   }
 

diff --git a/test/jest-test.ts b/test/jest-test.ts
@@ -34,13 +34,12 @@ describe("transform jest", () => {
       jest.mock('c', () => {}).mock('d', () => {});
       jest.doMock('a', () => {});
     `,
-      `
-jest.mock('a').unmock('b').enableAutomock().disableAutomock().mock('c', () => {}).mock('d', () => {});
+      `__jestHoist();__jestHoist2();__jestHoist3();__jestHoist4();__jestHoist5();__jestHoist6();
       import 'moduleName';
-;
-jest.unknown();
-
-;
+function __jestHoist(){jest.mock('a');};
+function __jestHoist2(){jest.unmock('b');}jest.unknown()function __jestHoist3(){jest.enableAutomock();};
+function __jestHoist4(){jest.disableAutomock();}
+function __jestHoist5(){jest.mock('c', () => {});}function __jestHoist6(){jest.mock('d', () => {});};
 jest.doMock('a', () => {});
     `,
     );
@@ -57,23 +56,21 @@ jest.doMock('a', () => {});
         jest.unmock('c')
       `,
       {
-        expectedCJSResult: `"use strict";${IMPORT_DEFAULT_PREFIX}
-jest.mock('a').mock('b', () => ({})).unmock('c');
+        expectedCJSResult: `"use strict";${IMPORT_DEFAULT_PREFIX}__jestHoist();__jestHoist2();__jestHoist3();
         var _a = require('a');
-;
+function __jestHoist(){jest.mock('a');};
         var _b = require('b');
-;
+function __jestHoist2(){jest.mock('b', () => ({}));};
         var _c = require('c'); var _c2 = _interopRequireDefault(_c);
-
+function __jestHoist3(){jest.unmock('c');}
       `,
-        expectedESMResult: `
-jest.mock('a').mock('b', () => ({})).unmock('c');
+        expectedESMResult: `__jestHoist();__jestHoist2();__jestHoist3();
         import {A} from 'a';
-;
+function __jestHoist(){jest.mock('a');};
         import {B} from 'b';
-;
+function __jestHoist2(){jest.mock('b', () => ({}));};
         import C from 'c';
-
+function __jestHoist3(){jest.unmock('c');}
       `,
       },
     );
@@ -88,11 +85,10 @@ jest.mock('a').mock('b', () => ({})).unmock('c');
 
       export const x = 1
     `,
-      `"use strict";${ESMODULE_PREFIX}${IMPORT_DEFAULT_PREFIX}
-jest.mock('a');
+      `"use strict";${ESMODULE_PREFIX}${IMPORT_DEFAULT_PREFIX}__jestHoist();
       var _a = require('./a'); var _a2 = _interopRequireDefault(_a);
       var _b = require('./b');
-;
+function __jestHoist(){jest.mock('a');};
 
        const x = 1; exports.x = x
     `,
@@ -109,18 +105,13 @@ jest.mock('a', () => ({
   }
 }));
     `,
-      `"use strict";${IMPORT_DEFAULT_PREFIX}${NULLISH_COALESCE_PREFIX}${OPTIONAL_CHAIN_PREFIX}
-jest.mock('a', () => ({
+      `"use strict";${IMPORT_DEFAULT_PREFIX}${NULLISH_COALESCE_PREFIX}${OPTIONAL_CHAIN_PREFIX}__jestHoist();
+      var _a = require('./a'); var _a2 = _interopRequireDefault(_a);
+function __jestHoist(){jest.mock('a', () => ({
   f(x) {
     return _nullishCoalesce(_optionalChain([x, 'optionalAccess', _ => _.a]), () => ( 0));
   }
-}));
-      var _a = require('./a'); var _a2 = _interopRequireDefault(_a);
-
-
-
-
-;
+}));};
     `,
     );
   });
@@ -136,18 +127,13 @@ jest.mock('a'! as number, (arg: unknown) => ({
 }) as any);
       x()
     `,
-      `"use strict";
-jest.mock('a' , (arg) => ({
+      `"use strict";__jestHoist();
+      var _a = require('./a');
+function __jestHoist(){jest.mock('a' , (arg) => ({
   f(x) {
     return x ;
   }
-}) );
-      var _a = require('./a');
-
-
-
-
-;
+}) );};
       _a.x.call(void 0, )
     `,
       {transforms: ["jsx", "jest", "imports", "typescript"]},
@@ -165,18 +151,13 @@ jest.mock('a': number, (arg: string) => ({
 }): any);
       x()
     `,
-      `"use strict";
-jest.mock('a', (arg) => ({
+      `"use strict";__jestHoist();
+      var _a = require('./a');
+function __jestHoist(){jest.mock('a', (arg) => ({
   f(x) {
     return (x);
   }
-}));
-      var _a = require('./a');
-
-
-
-
-;
+}));};
       _a.x.call(void 0, )
     `,
       {transforms: ["jsx", "jest", "imports", "flow"]},
@@ -195,19 +176,14 @@ jest.mock('a', (arg) => ({
       }));
       x()
     `,
-      `"use strict";${JSX_PREFIX}${IMPORT_DEFAULT_PREFIX}
-jest.mock('a', (arg) => ({
+      `"use strict";${JSX_PREFIX}${IMPORT_DEFAULT_PREFIX}__jestHoist();
+      var _react = require('react'); var _react2 = _interopRequireDefault(_react);
+      var _a = require('./a');
+function __jestHoist(){jest.mock('a', (arg) => ({
         f(x) {
           return _react2.default.createElement('div', {__self: this, __source: {fileName: _jsxFileName, lineNumber: 6}} );
         }
-      }));
-      var _react = require('react'); var _react2 = _interopRequireDefault(_react);
-      var _a = require('./a');
-
-
-
-
-;
+      }));};
       _a.x.call(void 0, )
     `,
       {transforms: ["jsx", "jest", "imports"]},
@@ -226,10 +202,9 @@ jest.mock('a', (arg) => ({
       _a.jest.mock('x');
     `,
         // Note that this behavior is incorrect, but jest requires imports transform for now.
-        expectedESMResult: `
-jest.mock('x');
+        expectedESMResult: `__jestHoist();
       import {jest} from './a';
-;
+function __jestHoist(){jest.mock('x');};
     `,
       },
     );
@@ -250,4 +225,18 @@ jest.mock('x');
       {transforms: ["jsx", "jest", "imports", "typescript"]},
     );
   });
+
+  it("allows chained unknown methods", () => {
+    assertResult(
+      `
+      import './a';
+      console.log(jest.spyOn({foo() {}}, 'foo').getMockName());
+    `,
+      `"use strict";
+      require('./a');
+      console.log(jest.spyOn({foo() {}}, 'foo').getMockName());
+    `,
+      {transforms: ["jest", "imports"]},
+    );
+  });
 });
diff --git a/website/package.json b/website/package.json
@@ -16,6 +16,7 @@
     "babel-core": "^6.26.3",
     "babel-jest": "^24.9.0",
     "babel-plugin-dynamic-import-node": "^2.3.0",
+    "babel-plugin-jest-hoist": "^26.6.2",
     "babel-runtime": "^6.26.0",
     "base64-js": "^1.3.1",
     "case-sensitive-paths-webpack-plugin": "^2.2.0",

diff --git a/website/src/Babel.ts b/website/src/Babel.ts
@@ -9,10 +9,13 @@ import {registerPlugin, transform} from "@babel/standalone";
 // @ts-ignore
 import DynamicImportPlugin from "babel-plugin-dynamic-import-node";
 // @ts-ignore
+import JestHoistPlugin from "babel-plugin-jest-hoist";
+// @ts-ignore
 import ReactHotLoaderPlugin from "react-hot-loader/dist/babel.development";
 
 registerPlugin("proposal-numeric-separator", NumericSeparatorPlugin);
 registerPlugin("dynamic-import-node", DynamicImportPlugin);
 registerPlugin("react-hot-loader", ReactHotLoaderPlugin);
+registerPlugin("jest-hoist", JestHoistPlugin);
 
 export {transform};
diff --git a/website/src/Constants.ts b/website/src/Constants.ts
@@ -42,6 +42,7 @@ export const TRANSFORMS: Array<TransformInfo> = [
   {name: "flow", presetName: "flow"},
   {name: "imports", babelName: "transform-modules-commonjs"},
   {name: "react-hot-loader", babelName: "react-hot-loader"},
+  {name: "jest", babelName: "jest-hoist"},
 ];
 
 export const DEFAULT_TRANSFORMS = ["jsx", "typescript", "imports"];