fix(prefer-expect-assertions): use scoped based jest fn call parser for expect checks

src/rules/__tests__/prefer-expect-assertions.test.ts

@@ -88,6 +88,20 @@ ruleTester.run('prefer-expect-assertions', rule, {
       `,
       options: [{ onlyFunctionsWithAsyncKeyword: true }],
     },
+    {
+      code: dedent`
+        import { expect as pleaseExpect } from '@jest/globals';
+
+        it("returns numbers that are greater than four", function() {
+          pleaseExpect.assertions(2);
+
+          for(let thing in things) {
+            pleaseExpect(number).toBeGreaterThan(4);
+          }
+        });
+      `,
+      parserOptions: { sourceType: 'module' },
+    },
   ],
   invalid: [
     {
@@ -120,11 +134,11 @@ ruleTester.run('prefer-expect-assertions', rule, {
           suggestions: [
             {
               messageId: 'suggestAddingHasAssertions',
-              output: 'it("it1", () => { expect.hasAssertions();foo()})',
+              output: 'it("it1", () => {expect.hasAssertions(); foo()})',
             },
             {
               messageId: 'suggestAddingAssertions',
-              output: 'it("it1", () => { expect.assertions();foo()})',
+              output: 'it("it1", () => {expect.assertions(); foo()})',
             },
           ],
         },
@@ -146,17 +160,17 @@ ruleTester.run('prefer-expect-assertions', rule, {
             {
               messageId: 'suggestAddingHasAssertions',
               output: dedent`
-                it("it1", function() {
-                  expect.hasAssertions();someFunctionToDo();
+                it("it1", function() {expect.hasAssertions();
+                  someFunctionToDo();
                   someFunctionToDo2();
                 });
               `,
             },
             {
               messageId: 'suggestAddingAssertions',
               output: dedent`
-                it("it1", function() {
-                  expect.assertions();someFunctionToDo();
+                it("it1", function() {expect.assertions();
+                  someFunctionToDo();
                   someFunctionToDo2();
                 });
               `,
@@ -1180,6 +1194,19 @@ ruleTester.run('prefer-expect-assertions (callbacks)', rule, {
         },
       ],
     },
+    {
+      code: dedent`
+        it("returns numbers that are greater than four", function(expect) {
+          expect.assertions(2);
+
+          for(let thing in things) {
+            expect(number).toBeGreaterThan(4);
+          }
+        });
+      `,
+      parserOptions: { sourceType: 'module' },
+      errors: [{ endColumn: 3, column: 1, messageId: 'haveExpectAssertions' }],
+    },
   ],
 });
 

src/rules/prefer-expect-assertions.ts

@@ -1,31 +1,29 @@
 import { AST_NODE_TYPES, TSESLint, TSESTree } from '@typescript-eslint/utils';
 import {
-  KnownCallExpression,
+  ParsedExpectFnCall,
   createRule,
   getAccessorValue,
-  hasOnlyOneArgument,
   isFunction,
-  isSupportedAccessor,
   isTypeOfJestFnCall,
   parseJestFnCall,
 } from './utils';
 
-const isExpectAssertionsOrHasAssertionsCall = (
-  expression: TSESTree.Node,
-): expression is KnownCallExpression<'assertions' | 'hasAssertions'> =>
-  expression.type === AST_NODE_TYPES.CallExpression &&
-  expression.callee.type === AST_NODE_TYPES.MemberExpression &&
-  isSupportedAccessor(expression.callee.object, 'expect') &&
-  isSupportedAccessor(expression.callee.property) &&
-  ['assertions', 'hasAssertions'].includes(
-    getAccessorValue(expression.callee.property),
-  );
+const isFirstStatement = (node: TSESTree.CallExpression): boolean => {
+  let parent: TSESTree.Node['parent'] = node;
+
+  while (parent) {
+    if (parent.parent?.type === AST_NODE_TYPES.BlockStatement) {
+      return parent.parent.body[0] === parent;
+    }
 
-const isFirstLineExprStmt = (
-  functionBody: TSESTree.Statement[],
-): functionBody is [TSESTree.ExpressionStatement] =>
-  functionBody[0] &&
-  functionBody[0].type === AST_NODE_TYPES.ExpressionStatement;
+    parent = parent.parent;
+  }
+
+  /* istanbul ignore next */
+  throw new Error(
+    `Could not find BlockStatement - please file a github issue at https://github.com/jest-community/eslint-plugin-jest`,
+  );
+};
 
 const suggestRemovingExtraArguments = (
   args: TSESTree.CallExpression['arguments'],
@@ -107,6 +105,7 @@ export default createRule<[RuleOptions], MessageIds>({
     let expressionDepth = 0;
     let hasExpectInCallback = false;
     let hasExpectInLoop = false;
+    let hasExpectAssertionsAsFirstStatement = false;
     let inTestCaseCall = false;
     let inForLoop = false;
 
@@ -140,6 +139,53 @@ export default createRule<[RuleOptions], MessageIds>({
       return false;
     };
 
+    const checkExpectHasAssertions = (expectFnCall: ParsedExpectFnCall) => {
+      if (getAccessorValue(expectFnCall.members[0]) === 'hasAssertions') {
+        if (expectFnCall.args.length) {
+          context.report({
+            messageId: 'hasAssertionsTakesNoArguments',
+            node: expectFnCall.matcher,
+            suggest: [suggestRemovingExtraArguments(expectFnCall.args, 0)],
+          });
+        }
+
+        return;
+      }
+
+      if (expectFnCall.args.length !== 1) {
+        let { loc } = expectFnCall.matcher;
+        const suggest: TSESLint.ReportSuggestionArray<MessageIds> = [];
+
+        if (expectFnCall.args.length) {
+          loc = expectFnCall.args[1].loc;
+          suggest.push(suggestRemovingExtraArguments(expectFnCall.args, 1));
+        }
+
+        context.report({
+          messageId: 'assertionsRequiresOneArgument',
+          suggest,
+          loc,
+        });
+
+        return;
+      }
+
+      const [arg] = expectFnCall.args;
+
+      if (
+        arg.type === AST_NODE_TYPES.Literal &&
+        typeof arg.value === 'number' &&
+        Number.isInteger(arg.value)
+      ) {
+        return;
+      }
+
+      context.report({
+        messageId: 'assertionsRequiresNumberArgument',
+        node: arg,
+      });
+    };
+
     const enterExpression = () => inTestCaseCall && expressionDepth++;
     const exitExpression = () => inTestCaseCall && expressionDepth--;
     const enterForLoop = () => (inForLoop = true);
@@ -166,6 +212,20 @@ export default createRule<[RuleOptions], MessageIds>({
         }
 
         if (jestFnCall?.type === 'expect' && inTestCaseCall) {
+          if (
+            expressionDepth === 1 &&
+            isFirstStatement(node) &&
+            jestFnCall.head.node.parent?.type ===
+              AST_NODE_TYPES.MemberExpression &&
+            jestFnCall.members.length === 1 &&
+            ['assertions', 'hasAssertions'].includes(
+              getAccessorValue(jestFnCall.members[0]),
+            )
+          ) {
+            checkExpectHasAssertions(jestFnCall);
+            hasExpectAssertionsAsFirstStatement = true;
+          }
+
           if (inForLoop) {
             hasExpectInLoop = true;
           }
@@ -202,92 +262,23 @@ export default createRule<[RuleOptions], MessageIds>({
         hasExpectInLoop = false;
         hasExpectInCallback = false;
 
-        const testFuncBody = testFn.body.body;
-
-        if (!isFirstLineExprStmt(testFuncBody)) {
-          context.report({
-            messageId: 'haveExpectAssertions',
-            node,
-            suggest: suggestions.map(([messageId, text]) => ({
-              messageId,
-              fix: fixer =>
-                fixer.insertTextBeforeRange(
-                  [testFn.body.range[0] + 1, testFn.body.range[1]],
-                  text,
-                ),
-            })),
-          });
+        if (hasExpectAssertionsAsFirstStatement) {
+          hasExpectAssertionsAsFirstStatement = false;
 
           return;
         }
 
-        const testFuncFirstLine = testFuncBody[0].expression;
-
-        if (!isExpectAssertionsOrHasAssertionsCall(testFuncFirstLine)) {
-          context.report({
-            messageId: 'haveExpectAssertions',
-            node,
-            suggest: suggestions.map(([messageId, text]) => ({
-              messageId,
-              fix: fixer => fixer.insertTextBefore(testFuncBody[0], text),
-            })),
-          });
-
-          return;
-        }
-
-        if (
-          isSupportedAccessor(
-            testFuncFirstLine.callee.property,
-            'hasAssertions',
-          )
-        ) {
-          if (testFuncFirstLine.arguments.length) {
-            context.report({
-              messageId: 'hasAssertionsTakesNoArguments',
-              node: testFuncFirstLine.callee.property,
-              suggest: [
-                suggestRemovingExtraArguments(testFuncFirstLine.arguments, 0),
-              ],
-            });
-          }
-
-          return;
-        }
-
-        if (!hasOnlyOneArgument(testFuncFirstLine)) {
-          let { loc } = testFuncFirstLine.callee.property;
-          const suggest: TSESLint.ReportSuggestionArray<MessageIds> = [];
-
-          if (testFuncFirstLine.arguments.length) {
-            loc = testFuncFirstLine.arguments[1].loc;
-            suggest.push(
-              suggestRemovingExtraArguments(testFuncFirstLine.arguments, 1),
-            );
-          }
-
-          context.report({
-            messageId: 'assertionsRequiresOneArgument',
-            suggest,
-            loc,
-          });
-
-          return;
-        }
-
-        const [arg] = testFuncFirstLine.arguments;
-
-        if (
-          arg.type === AST_NODE_TYPES.Literal &&
-          typeof arg.value === 'number' &&
-          Number.isInteger(arg.value)
-        ) {
-          return;
-        }
-
         context.report({
-          messageId: 'assertionsRequiresNumberArgument',
-          node: arg,
+          messageId: 'haveExpectAssertions',
+          node,
+          suggest: suggestions.map(([messageId, text]) => ({
+            messageId,
+            fix: fixer =>
+              fixer.insertTextBeforeRange(
+                [testFn.body.range[0] + 1, testFn.body.range[1]],
+                text,
+              ),
+          })),
         });
       },
     };