write some codemods for jest.mock

package.json

@@ -6,7 +6,7 @@
     "format:check": "prettier . --check",
     "format:fix": "prettier . --write",
     "lint:check": "eslint",
-    "lint:fix": "eslint --fix",
+    "lint:fix": "eslint --fix --no-warn-ignored",
     "lint:staged": "lint-staged",
     "prepare": "husky",
     "test": "jest",

package/plugin/codemods/jest-mock/extract_mocked_specifier.js

@@ -0,0 +1,26 @@
+const babelTypes = require("@babel/types");
+
+function extract_mocked_specifier(path, specifier, callback) {
+  // console.log(path.parent)
+  const expression = path.parent;
+
+  if (path.node.arguments[1].body.properties.length === 1) {
+    callback?.(path.parent);
+    return;
+  }
+
+  const clone = babelTypes.cloneNode(expression);
+  clone.expression.arguments[1].body.properties =
+    clone.expression.arguments[1].body.properties.filter((prop) => {
+      return prop.key.name === specifier;
+    });
+  path.node.arguments[1].body.properties =
+    path.node.arguments[1].body.properties.filter((prop) => {
+      return prop.key.name !== specifier;
+    });
+  callback?.(clone);
+
+  path.insertBefore(clone);
+}
+
+module.exports = { extract_mocked_specifier };

package/plugin/codemods/jest-mock/extract_mocked_specifier.test.js

@@ -0,0 +1,122 @@
+const { babelParse } = require("../../utils");
+const traverse = require("@babel/traverse").default;
+const babelGenerate = require("@babel/generator").default;
+const { extract_mocked_specifier } = require("./extract_mocked_specifier");
+const {
+  multilineTrim,
+} = require("../../../../spec/test-utils/expectTransform");
+
+function parse(code, callback) {
+  const ast = babelParse(code);
+
+  traverse(ast, {
+    CallExpression(path) {
+      if (path.node.callee?.property?.name === "mock") {
+        if (path.node.arguments.length > 1) {
+          callback(path);
+        }
+      }
+    },
+  });
+
+  return ast;
+}
+
+function expectCodemod(code, codemod, expectedOutput) {
+  const ast = parse(code, (path) => {
+    return codemod(path);
+  });
+
+  const output = multilineTrim(babelGenerate(ast).code);
+
+  expect(output).toBe(multilineTrim(expectedOutput));
+}
+
+describe("codemods/jest-mock/extract_mocked_specifier", () => {
+  it("leaves code as is if there is only one key", () => {
+    expectCodemod(
+      `
+      jest.mock('./origin.js', () => ({
+        a: 'a'
+      }));
+      `,
+      (path) => extract_mocked_specifier(path, "a"),
+      `
+      jest.mock('./origin.js', () => ({
+        a: 'a'
+      }));
+      `,
+    );
+  });
+
+  it("extracts a separate jest.mock statement for a given specifier", () => {
+    expectCodemod(
+      `
+      jest.mock('./origin.js', () => ({
+        a: 'a',
+        b: 'b'
+      }));
+      `,
+      (path) => extract_mocked_specifier(path, "b"),
+      `
+      jest.mock('./origin.js', () => ({
+        b: 'b'
+      }));
+      jest.mock('./origin.js', () => ({
+        a: 'a'
+      }));
+      `,
+    );
+  });
+
+  it("can extract multiple specifiers", () => {
+    expectCodemod(
+      `
+      jest.mock('./origin.js', () => ({
+        a: 'a',
+        b: 'b',
+        c: 'c'
+      }));
+      `,
+      (path) => {
+        extract_mocked_specifier(path, "b");
+        extract_mocked_specifier(path, "c");
+      },
+      `
+      jest.mock('./origin.js', () => ({
+        b: 'b'
+      }));
+      jest.mock('./origin.js', () => ({
+        c: 'c'
+      }));
+      jest.mock('./origin.js', () => ({
+        a: 'a'
+      }));
+      `,
+    );
+  });
+
+  it("returns the created node for further processing", () => {
+    expectCodemod(
+      `
+      jest.mock('./origin.js', () => ({
+        a: 'a',
+        b: 'b'
+      }));
+      `,
+      (path) => {
+        extract_mocked_specifier(path, "b", (node) => {
+          node.expression.arguments[1].body.properties[0].key.name = "bar";
+        });
+      },
+      `
+      jest.mock('./origin.js', () => ({
+        bar: 'b'
+      }));
+      jest.mock('./origin.js', () => ({
+        a: 'a'
+      }));
+      `,
+    );
+  });
+});

package/plugin/index.js

@@ -3,6 +3,9 @@ const { resolve } = require("../resolve");
 const { withCache } = require("../cache");
 const { matchAnyRegex, removeItemsByIndexesInPlace } = require("./utils");
 const { Tracer } = require("./traverse");
+const {
+  extract_mocked_specifier,
+} = require("./codemods/jest-mock/extract_mocked_specifier");
 
 let modulePaths = null;
 let moduleNameMapper = null;
@@ -55,15 +58,47 @@ module.exports = function babelPlugin(babel) {
           return;
         }
         if (path.node.callee?.property?.name === "mock") {
-          const mockedPath = path.node.arguments[0].value;
           const resolved = bjbResolve(
-            mockedPath,
+            path.node.arguments[0].value,
             nodepath.dirname(state.file.opts.filename),
           );
           if (isPathWhitelisted(resolved)) {
             return;
           }
+
           path.node.arguments[0].value = resolved;
+
+          const importedFrom = resolved;
+
+          if (path.node.arguments.length > 1) {
+            // jest.mock('./origin.js', () => ({ target: value }))
+            // Figure out the mocked specifier, trace it and re-write the mock call;
+            try {
+              path.node.arguments?.[1]?.body?.properties?.forEach((objProp) => {
+                const mockedIdentifier = objProp?.key?.name;
+
+                const specifierOrigin = traceSpecifierOrigin(
+                  mockedIdentifier,
+                  importedFrom,
+                );
+
+                if (!specifierOrigin) {
+                  return;
+                }
+
+                if (importedFrom === specifierOrigin.source) {
+                  // specifier is imported from the already resolved place, skip
+                  return;
+                }
+
+                extract_mocked_specifier(path, mockedIdentifier, (node) => {
+                  node.expression.arguments[0].value = specifierOrigin.source;
+                });
+              });
+            } catch (e) {
+              console.log("failed to parse mock statement: ", e);
+            }
+          }
         }
       },
       ImportDeclaration(path, state) {

spec/import-cases.spec.js → spec/import-rewrites.spec.jsm

@@ -8,7 +8,7 @@ const { createExpectTransform } = require("./test-utils/expectTransform.js");
 
 const expectTransform = createExpectTransform(__filename);
 
-describe("babel-jest-boost plugin import cases", () => {
+describe("babel-jest-boost plugin import rewrites", () => {
   it("correctly traces specifiers within all import syntaxes", () => {
     expectTransform(
       "import target from './test_tree/default';",

spec/jest-mock-rewrites.spec.js

@@ -0,0 +1,61 @@
+const {
+  createExpectTransform,
+  createExpectJestTransform,
+} = require("./test-utils/expectTransform.js");
+
+const expectJestTransform = createExpectJestTransform(__filename);
+const expectTransform = createExpectTransform(__filename);
+
+describe("babel-jest-boost plugin jest.mock rewrites", () => {
+  it("resolves modules mocked in jest.mock", () => {
+    expectTransform(
+      "import { a } from './test_tree/consts'",
+      `import { a } from "${__dirname}/test_tree/consts/a.js";`,
+    );
+    expectJestTransform(
+      "jest.mock('./test_tree/default');",
+      `_getJestObj().mock("${__dirname}/test_tree/default/index.js");`,
+    );
+
+    expectJestTransform(
+      "jest.mock('./test_tree/default', () => {});",
+      `_getJestObj().mock("${__dirname}/test_tree/default/index.js", () => {});`,
+    );
+
+    expectJestTransform(
+      `jest.mock('./test_tree/consts', () => ({ a: 1 }));`,
+      `
+        _getJestObj().mock("${__dirname}/test_tree/consts/a.js", () => ({
+          a: 1
+        }));
+      `,
+    );
+    expectJestTransform(
+      `jest.mock('./test_tree/consts', () => ({ a: 1, b: 2 }));`,
+      `
+        _getJestObj().mock("${__dirname}/test_tree/consts/b.js", () => ({
+          b: 2
+        }));
+        _getJestObj().mock("${__dirname}/test_tree/consts/a.js", () => ({
+          a: 1
+        }));
+      `,
+    );
+
+    expectJestTransform(
+      `jest.mock("${__dirname}/test_tree/consts/consts.js", () => ({ a: 1, b: 2, c: 3, d: 4 }));`,
+      `
+        _getJestObj().mock("${__dirname}/test_tree/consts/consts.js", () => ({
+          c: 3,
+          d: 4
+        }));
+        _getJestObj().mock("${__dirname}/test_tree/consts/a.js", () => ({
+          a: 1
+        }));
+        _getJestObj().mock("${__dirname}/test_tree/consts/b.js", () => ({
+          b: 2
+        }));
+      `,
+    );
+  });
+});

spec/test-utils/expectTransform.js

@@ -26,6 +26,28 @@ function multilineTrim(string) {
     .join("\n");
 }
 
+function multilineRemove(mainString, stringToRemove) {
+  return mainString.replace(
+    new RegExp(stringToRemove.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"), "g"),
+    "",
+  );
+}
+
+function removeJestBlob(input) {
+  const jestBlob =
+    "\n" +
+    multilineTrim(`
+    function _getJestObj() {
+      const {
+      jest
+      } = require("@jest/globals");
+      _getJestObj = () => jest;
+      return jest;
+    }
+  `);
+
+  return multilineRemove(input, jestBlob);
+}
 function createExpectTransform(filename, options) {
   const transformer = createTransform(options);
 
@@ -39,4 +61,24 @@ function createExpectTransform(filename, options) {
   };
 }
 
-module.exports = { createExpectTransform };
+function createExpectJestTransform(filename, options) {
+  const transformer = createTransform(options);
+
+  function transform(source) {
+    const output = transformer.process(source, filename, { config: {} });
+    return multilineTrim(output.code.replace(/\/\/# sourceMappingURL.*/, ""));
+  }
+
+  return function expectJestTransform(input, expectedOutput) {
+    expectedOutput = multilineTrim(expectedOutput);
+    const output = removeJestBlob(transform(input));
+
+    expect(output).toBe(multilineTrim(expectedOutput));
+  };
+}
+
+module.exports = {
+  createExpectTransform,
+  createExpectJestTransform,
+  multilineTrim,
+};

spec/test_tree/consts/a.js

@@ -0,0 +1 @@
+export const a = 1;

spec/test_tree/consts/b.js

@@ -0,0 +1 @@
+export const b = 2;

spec/test_tree/consts/consts.js

@@ -0,0 +1,4 @@
+export * from "./a.js";
+export * from "./b.js";
+export const c = 3;
+export const d = 4;

spec/test_tree/consts/index.js

@@ -0,0 +1 @@
+export * from "./consts";