From d0cea37ae0a8ab07b8082cedbaaf161bcc94c405 Mon Sep 17 00:00:00 2001
From: Mario Campa <mcampa@gmail.com>
Date: Mon, 5 Oct 2020 15:15:05 -0700
Subject: [PATCH] feat(prefer-expect-assertions): add
 `onlyFunctionsWithAsyncKeyword` option (#677)

Co-authored-by: Mario Campa <mcampa@twitter.com>
---
 docs/rules/prefer-expect-assertions.md        | 42 +++++++++++++++++++
 .../prefer-expect-assertions.test.ts          | 39 ++++++++++++++++-
 src/rules/prefer-expect-assertions.ts         | 25 +++++++++--
 3 files changed, 101 insertions(+), 5 deletions(-)

diff --git a/docs/rules/prefer-expect-assertions.md b/docs/rules/prefer-expect-assertions.md
index d99720356..b80b63dc3 100644
--- a/docs/rules/prefer-expect-assertions.md
+++ b/docs/rules/prefer-expect-assertions.md
@@ -55,3 +55,45 @@ test('my test', () => {
   expect(someThing()).toEqual('foo');
 });
 ```
+
+## Options
+
+#### `onlyFunctionsWithAsyncKeyword`
+
+When `true`, this rule will only warn for tests that use the `async` keyword.
+
+```json
+{
+  "rules": {
+    "jest/prefer-expect-assertions": [
+      "warn",
+      { "onlyFunctionsWithAsyncKeyword": true }
+    ]
+  }
+}
+```
+
+When `onlyFunctionsWithAsyncKeyword` option is set to `true`, the following
+pattern would be a warning:
+
+```js
+test('my test', async () => {
+  const result = await someAsyncFunc();
+  expect(result).toBe('foo');
+});
+```
+
+While the following patterns would not be considered warnings:
+
+```js
+test('my test', () => {
+  const result = someFunction();
+  expect(result).toBe('foo');
+});
+
+test('my test', async () => {
+  expect.assertions(1);
+  const result = await someAsyncFunc();
+  expect(result).toBe('foo');
+});
+```
diff --git a/src/rules/__tests__/prefer-expect-assertions.test.ts b/src/rules/__tests__/prefer-expect-assertions.test.ts
index 648b7904f..8f7e57d88 100644
--- a/src/rules/__tests__/prefer-expect-assertions.test.ts
+++ b/src/rules/__tests__/prefer-expect-assertions.test.ts
@@ -6,7 +6,7 @@ import rule from '../prefer-expect-assertions';
 const ruleTester = new TSESLint.RuleTester({
   parser: resolveFrom(require.resolve('eslint'), 'espree'),
   parserOptions: {
-    ecmaVersion: 2015,
+    ecmaVersion: 2017,
   },
 });
 
@@ -205,6 +205,21 @@ ruleTester.run('prefer-expect-assertions', rule, {
         },
       ],
     },
+    {
+      code: dedent`
+        it("it1", async function() {
+          expect(someValue).toBe(true);
+        })
+      `,
+      options: [{ onlyFunctionsWithAsyncKeyword: true }],
+      errors: [
+        {
+          messageId: 'haveExpectAssertions',
+          column: 1,
+          line: 1,
+        },
+      ],
+    },
   ],
 
   valid: [
@@ -223,5 +238,27 @@ ruleTester.run('prefer-expect-assertions', rule, {
     'test("it1")',
     'itHappensToStartWithIt("foo", function() {})',
     'testSomething("bar", function() {})',
+    'it(async () => {expect.assertions(0);})',
+    {
+      code: dedent`
+        it("it1", async () => {
+          expect.assertions(1);
+          expect(someValue).toBe(true)
+        })
+      `,
+      options: [{ onlyFunctionsWithAsyncKeyword: true }],
+    },
+    {
+      code: dedent`
+        it("it1", function() {
+          expect(someValue).toBe(true)
+        })
+      `,
+      options: [{ onlyFunctionsWithAsyncKeyword: true }],
+    },
+    {
+      code: 'it("it1", () => {})',
+      options: [{ onlyFunctionsWithAsyncKeyword: true }],
+    },
   ],
 });
diff --git a/src/rules/prefer-expect-assertions.ts b/src/rules/prefer-expect-assertions.ts
index 6cc707da9..c65ad3c6b 100644
--- a/src/rules/prefer-expect-assertions.ts
+++ b/src/rules/prefer-expect-assertions.ts
@@ -46,6 +46,9 @@ interface PreferExpectAssertionsCallExpression extends TSESTree.CallExpression {
     TSESTree.ArrowFunctionExpression & { body: TSESTree.BlockStatement },
   ];
 }
+interface RuleOptions {
+  onlyFunctionsWithAsyncKeyword?: boolean;
+}
 
 type MessageIds =
   | 'hasAssertionsTakesNoArguments'
@@ -61,7 +64,7 @@ const suggestions: Array<[MessageIds, string]> = [
   ['suggestAddingAssertions', 'expect.assertions();'],
 ];
 
-export default createRule<[], MessageIds>({
+export default createRule<[RuleOptions], MessageIds>({
   name: __filename,
   meta: {
     docs: {
@@ -85,14 +88,28 @@ export default createRule<[], MessageIds>({
       suggestRemovingExtraArguments: 'Remove extra arguments',
     },
     type: 'suggestion',
-    schema: [],
+    schema: [
+      {
+        type: 'object',
+        properties: {
+          onlyFunctionsWithAsyncKeyword: {
+            type: 'boolean',
+          },
+        },
+        additionalProperties: false,
+      },
+    ],
   },
-  defaultOptions: [],
-  create(context) {
+  defaultOptions: [{ onlyFunctionsWithAsyncKeyword: false }],
+  create(context, [options]) {
     return {
       'CallExpression[callee.name=/^(it|test)$/][arguments.1.body.body]'(
         node: PreferExpectAssertionsCallExpression,
       ) {
+        if (options.onlyFunctionsWithAsyncKeyword && !node.arguments[1].async) {
+          return;
+        }
+
         const testFuncBody = node.arguments[1].body.body;
 
         if (!isFirstLineExprStmt(testFuncBody)) {