From b3303b7d4c1bc6d11b17159d3a6b998707bb12f6 Mon Sep 17 00:00:00 2001
From: Yukihiro Hasegawa <49516827+y-hsgw@users.noreply.github.com>
Date: Wed, 18 Dec 2024 08:08:31 +0900
Subject: [PATCH] [prefer-expect-assertions]fix false negatives for test
 context usage (#600)

* fix: fix false negatives for test context usage

* test: add test
---
 src/rules/prefer-expect-assertions.ts  | 10 +++++--
 src/utils/parse-vitest-fn-call.ts      |  6 +++-
 tests/prefer-expect-assertions.test.ts | 41 ++++++++++++++++++++++++++
 3 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/src/rules/prefer-expect-assertions.ts b/src/rules/prefer-expect-assertions.ts
index d6df45d..1b7de9a 100644
--- a/src/rules/prefer-expect-assertions.ts
+++ b/src/rules/prefer-expect-assertions.ts
@@ -88,6 +88,7 @@ export default createEslintRule<Options[], MessageIds>({
     let hasExpectInCallBack = false
     let hasExpectInLoop = false
     let hasExpectAssertAsFirstStatement = false
+    let testContextName: string | null = null
     let inTestCaseCall = false
     let inForLoop = false
 
@@ -176,6 +177,10 @@ export default createEslintRule<Options[], MessageIds>({
           return
         }
 
+        if (vitestFnCall?.head.type === "testContext" && vitestFnCall.members[0].type === AST_NODE_TYPES.Identifier && vitestFnCall.members[0].name === "expect") {
+          testContextName = `${vitestFnCall.head.local}`
+        }
+
         if (vitestFnCall?.type === 'expect' && inTestCaseCall) {
           if (expressionDepth === 1 && isFirstStatement(node) && vitestFnCall.head.node.parent?.type === AST_NODE_TYPES.MemberExpression && vitestFnCall.members.length === 1
             && ['assertions', 'hasAssertions'].includes(getAccessorValue(vitestFnCall.members[0]))) {
@@ -217,8 +222,9 @@ export default createEslintRule<Options[], MessageIds>({
         const suggestions: Array<[MessageIds, string]> = []
 
         if (secondArg.body.type === AST_NODE_TYPES.BlockStatement) {
-          suggestions.push(['suggestAddingHasAssertions', 'expect.hasAssertions();'],
-            ['suggestAddingAssertions', 'expect.assertions();'])
+          const prefix = testContextName ? `${testContextName}.` : "";
+          suggestions.push(['suggestAddingHasAssertions', `${prefix}expect.hasAssertions();`],
+            ['suggestAddingAssertions', `${prefix}expect.assertions();`])
         }
 
         context.report({
diff --git a/src/utils/parse-vitest-fn-call.ts b/src/utils/parse-vitest-fn-call.ts
index abc5e6d..5b15aa8 100644
--- a/src/utils/parse-vitest-fn-call.ts
+++ b/src/utils/parse-vitest-fn-call.ts
@@ -252,7 +252,7 @@ const parseVitestFnCallWithReasonInner = (
 
   const links = [name, ...rest.map(getAccessorValue)]
 
-  if (name !== 'vi' && name !== 'expect' && name !== 'expectTypeOf' && !ValidVitestFnCallChains.has(links.join('.')))
+  if (resolved.type !== "testContext" && name !== 'vi' && name !== 'expect' && name !== 'expectTypeOf' && !ValidVitestFnCallChains.has(links.join('.')))
     return null
 
   const parsedVitestFnCall: Omit<ParsedVitestFnCall, 'type'> = {
@@ -382,6 +382,10 @@ export const resolveScope = (
           return "testContext"
       }
 
+      const namedParam = isFunction(def.node) ? def.node.params.find(params => params.type === AST_NODE_TYPES.Identifier ) : undefined
+      if (namedParam)
+        return "testContext"
+
       const importDetails = describePossibleImportDef(def)
 
       if (importDetails?.local === identifier)
diff --git a/tests/prefer-expect-assertions.test.ts b/tests/prefer-expect-assertions.test.ts
index b7ca47c..eaf7266 100644
--- a/tests/prefer-expect-assertions.test.ts
+++ b/tests/prefer-expect-assertions.test.ts
@@ -98,6 +98,47 @@ it('my test description', ({ expect }) => {expect.assertions();
 
   expect(sum(a, b)).toBe(a + b);
 })
+`
+            }
+          ]
+        }
+      ],
+    },
+    {
+      code: `
+it('my test description', (context) => {
+  const a = 1;
+  const b = 2;
+
+  context.expect(sum(a, b)).toBe(a + b);
+})
+`,
+      errors: [
+        {
+          messageId: 'haveExpectAssertions',
+          column: 1,
+          line: 2,
+          suggestions: [
+            {
+              messageId: 'suggestAddingHasAssertions',
+              output: `
+it('my test description', (context) => {context.expect.hasAssertions();
+  const a = 1;
+  const b = 2;
+
+  context.expect(sum(a, b)).toBe(a + b);
+})
+`
+            },
+            {
+              messageId: 'suggestAddingAssertions',
+              output: `
+it('my test description', (context) => {context.expect.assertions();
+  const a = 1;
+  const b = 2;
+
+  context.expect(sum(a, b)).toBe(a + b);
+})
 `
             }
           ]