Codemod for auto-creating @jest/globals imports (#508)

* Codemod for auto-creating @jest/globals imports

* use utils/imports

* Update README

* hardcode @jest/globals API + add test

* Fix lockfile

---------

Co-authored-by: skovhus <kenneth.skovhus@gmail.com>

README.md

@@ -84,6 +84,7 @@ $ jscodeshift -t node_modules/jest-codemods/dist/transformers/mocha.js test-fold
 $ jscodeshift -t node_modules/jest-codemods/dist/transformers/should.js test-folder
 $ jscodeshift -t node_modules/jest-codemods/dist/transformers/tape.js test-folder
 $ jscodeshift -t node_modules/jest-codemods/dist/transformers/sinon.js test-folder
+$ jscodeshift -t node_modules/jest-codemods/dist/transformers/jest-globals-import.js test-folder
 ```
 
 ## Test environment: Jest on Node.js or other

package.json

@@ -63,6 +63,7 @@
     "update-notifier": "5.1.0"
   },
   "devDependencies": {
+    "@jest/globals": "29.3.1",
     "@types/jest": "29.2.6",
     "@types/jscodeshift": "0.11.6",
     "@types/update-notifier": "5.1.0",

src/transformers/jest-globals-import.test.ts

@@ -0,0 +1,144 @@
+/* eslint-env jest */
+import fs from 'fs'
+import * as jscodeshift from 'jscodeshift'
+import { applyTransform } from 'jscodeshift/src/testUtils'
+import path from 'path'
+
+import { JEST_GLOBALS } from '../utils/consts'
+import * as plugin from './jest-globals-import'
+
+function expectTransformation(source: string, expectedOutput: string | null) {
+  const result = applyTransform({ ...plugin, parser: 'ts' }, {}, { source })
+  expect(result).toBe(expectedOutput ?? '')
+}
+
+describe('jestGlobalsImport', () => {
+  it("matches @jest/globals' types", () => {
+    const jestGlobalsPath = path.join(
+      __dirname,
+      '../../node_modules/@jest/globals/build/index.d.ts'
+    )
+
+    const jestGlobals = new Set<string>()
+
+    const j = jscodeshift.withParser('ts')
+    const jestGlobalsAst = j(String(fs.readFileSync(jestGlobalsPath)))
+
+    jestGlobalsAst
+      .find(j.ExportNamedDeclaration, { declaration: { declare: true } })
+      .forEach((exportNamedDec) => {
+        if (exportNamedDec.node.declaration?.type !== 'VariableDeclaration') return
+        exportNamedDec.node.declaration.declarations.forEach((dec) => {
+          if (dec.type !== 'VariableDeclarator' || dec.id?.type !== 'Identifier') return
+          jestGlobals.add(dec.id.name)
+        })
+      })
+
+    jestGlobalsAst
+      .find(j.ExportSpecifier, { exported: { name: (n) => typeof n === 'string' } })
+      .forEach((exportSpecifier) => {
+        jestGlobals.add(exportSpecifier.node.exported.name)
+      })
+
+    expect(jestGlobals).toEqual(JEST_GLOBALS)
+  })
+
+  it('covers a simple test', () => {
+    expectTransformation(
+      `
+it('works', () => {
+  expect(true).toBe(true);
+});
+`.trim(),
+      `
+import { expect, it } from '@jest/globals';
+it('works', () => {
+  expect(true).toBe(true);
+});
+`.trim()
+    )
+  })
+
+  it('ignores locally defined variables with the same name', () => {
+    expectTransformation(
+      `
+const test = () => { console.log('only a test'); };
+{
+  function b() {
+    function c() {
+      test();
+    }
+  }
+}
+`.trim(),
+      null
+    )
+  })
+
+  it('removes imports', () => {
+    expectTransformation(
+      `
+import '@jest/globals';
+const BLAH = 5;
+`.trim(),
+      `
+const BLAH = 5;
+`.trim()
+    )
+    expectTransformation(
+      `
+import { expect } from '@jest/globals';
+const BLAH = 5;
+`.trim(),
+      `
+const BLAH = 5;
+`.trim()
+    )
+    expectTransformation(
+      `
+import * as jestGlobals from '@jest/globals';
+const BLAH = 5;
+`.trim(),
+      `
+const BLAH = 5;
+`.trim()
+    )
+  })
+
+  it('covers a less simple test', () => {
+    expectTransformation(
+      `
+import { expect as xpect, it } from '@jest/globals';
+import wrapWithStuff from 'test-utils/wrapWithStuff';
+
+describe('with foo=bar', () => {
+  wrapWithStuff({ foo: 'bar' });
+
+  beforeEach(() => jest.useFakeTimers());
+  afterEach(() => jest.useRealTimers());
+
+  it('works', () => {
+    xpect(myThingIsEnabled(jest.fn())).toBe(true);
+    expect(1).toBe(1);
+  });
+});
+`.trim(),
+      `
+import { expect as xpect, it, afterEach, beforeEach, describe, jest } from '@jest/globals';
+import wrapWithStuff from 'test-utils/wrapWithStuff';
+
+describe('with foo=bar', () => {
+  wrapWithStuff({ foo: 'bar' });
+
+  beforeEach(() => jest.useFakeTimers());
+  afterEach(() => jest.useRealTimers());
+
+  it('works', () => {
+    xpect(myThingIsEnabled(jest.fn())).toBe(true);
+    expect(1).toBe(1);
+  });
+});
+`.trim()
+    )
+  })
+})

src/transformers/jest-globals-import.ts

@@ -0,0 +1,58 @@
+import type { ImportSpecifier, JSCodeshift } from 'jscodeshift'
+
+import { JEST_GLOBALS } from '../utils/consts'
+import { findImports, removeRequireAndImport } from '../utils/imports'
+
+const jestGlobalsImport = (
+  fileInfo: { path: string; source: string },
+  api: { jscodeshift: JSCodeshift }
+) => {
+  const { jscodeshift: j } = api
+  const ast = j(fileInfo.source)
+
+  const jestGlobalsUsed = new Set<string>()
+
+  JEST_GLOBALS.forEach((globalName) => {
+    if (
+      ast
+        .find(j.CallExpression, { callee: { name: globalName } })
+        .filter((callExpression) => !callExpression.scope.lookup(globalName))
+        .size() > 0 ||
+      ast
+        .find(j.MemberExpression, { object: { name: globalName } })
+        .filter((memberExpression) => !memberExpression.scope.lookup(globalName))
+        .size() > 0
+    ) {
+      jestGlobalsUsed.add(globalName)
+    }
+  })
+
+  const jestGlobalsImports = findImports(j, ast, '@jest/globals')
+  const hasJestGlobalsImport = jestGlobalsImports.length > 0
+  const needsJestGlobalsImport = jestGlobalsUsed.size > 0
+
+  if (!needsJestGlobalsImport) {
+    if (!hasJestGlobalsImport) return null
+    removeRequireAndImport(j, ast, '@jest/globals')
+  } else {
+    const jestGlobalsImport = hasJestGlobalsImport
+      ? jestGlobalsImports.get().value
+      : j.importDeclaration([], j.stringLiteral('@jest/globals'))
+    const { specifiers } = jestGlobalsImport
+    const existingNames = new Set<string>(
+      specifiers.map((s: ImportSpecifier) => s.imported.name)
+    )
+    jestGlobalsUsed.forEach((jestGlobal) => {
+      if (!existingNames.has(jestGlobal)) {
+        specifiers.push(j.importSpecifier(j.identifier(jestGlobal)))
+      }
+    })
+    if (!hasJestGlobalsImport) {
+      ast.find(j.Program).get('body', 0).insertBefore(jestGlobalsImport)
+    }
+  }
+
+  return ast.toSource({ quote: 'single' })
+}
+
+export default jestGlobalsImport

src/utils/consts.ts

@@ -44,3 +44,20 @@ export const JEST_MATCHER_TO_MAX_ARGS = {
 }
 
 export const JEST_MOCK_PROPERTIES = new Set(['spyOn', 'fn', 'createSpy'])
+
+export const JEST_GLOBALS = new Set<string>([
+  'afterAll',
+  'afterEach',
+  'beforeAll',
+  'beforeEach',
+  'describe',
+  'expect',
+  'fdescribe',
+  'fit',
+  'it',
+  'jest',
+  'test',
+  'xdescribe',
+  'xit',
+  'xtest',
+])