From bd2c33c858573d5414d8bc0d401eb6f27801ad2b Mon Sep 17 00:00:00 2001
From: Gareth Jones <Jones258@Gmail.com>
Date: Tue, 12 Oct 2021 08:50:24 +1300
Subject: [PATCH] fix(valid-expect-in-promise): support awaited promises in
 arguments (#936)

---
 .../__tests__/valid-expect-in-promise.test.ts | 100 ++++++++++++++++++
 src/rules/valid-expect-in-promise.ts          |  42 ++++++++
 2 files changed, 142 insertions(+)

diff --git a/src/rules/__tests__/valid-expect-in-promise.test.ts b/src/rules/__tests__/valid-expect-in-promise.test.ts
index 7a9a4ccf6..51ae14ff4 100644
--- a/src/rules/__tests__/valid-expect-in-promise.test.ts
+++ b/src/rules/__tests__/valid-expect-in-promise.test.ts
@@ -15,6 +15,88 @@ ruleTester.run('valid-expect-in-promise', rule, {
     "test('something', () => Promise.resolve().then(() => expect(1).toBe(2)));",
     'Promise.resolve().then(() => expect(1).toBe(2))',
     'const x = Promise.resolve().then(() => expect(1).toBe(2))',
+    dedent`
+      it('is valid', async () => {
+        const promise = loadNumber().then(number => {
+          expect(typeof number).toBe('number');
+
+          return number + 1;
+        });
+
+        expect(await promise).toBeGreaterThan(1);
+      });
+    `,
+    dedent`
+      it('is valid', async () => {
+        const promise = loadNumber().then(number => {
+          expect(typeof number).toBe('number');
+
+          return number + 1;
+        });
+
+        expect(await promise).resolves.toBeGreaterThan(1);
+      });
+    `,
+    dedent`
+      it('is valid', async () => {
+        const promise = loadNumber().then(number => {
+          expect(typeof number).toBe('number');
+
+          return number + 1;
+        });
+
+        expect(1).toBeGreaterThan(await promise);
+      });
+    `,
+    dedent`
+      it('is valid', async () => {
+        const promise = loadNumber().then(number => {
+          expect(typeof number).toBe('number');
+
+          return number + 1;
+        });
+
+        expect.this.that.is(await promise);
+      });
+    `,
+    dedent`
+      it('is valid', async () => {
+        expect(await loadNumber().then(number => {
+          expect(typeof number).toBe('number');
+
+          return number + 1;
+        })).toBeGreaterThan(1);
+      });
+    `,
+    dedent`
+      it('is valid', async () => {
+        const promise = loadNumber().then(number => {
+          expect(typeof number).toBe('number');
+
+          return number + 1;
+        });
+
+        logValue(await promise);
+      });
+    `,
+    dedent`
+      it('is valid', async () => {
+        const promise = loadNumber().then(number => {
+          expect(typeof number).toBe('number');
+
+          return 1;
+        });
+
+        expect.assertions(await promise);
+      });
+    `,
+    dedent`
+      it('is valid', async () => {
+        await loadNumber().then(number => {
+          expect(typeof number).toBe('number');
+        });
+      });
+    `,
     dedent`
       it('it1', () => new Promise((done) => {
         test()
@@ -1369,5 +1451,23 @@ ruleTester.run('valid-expect-in-promise', rule, {
         },
       ],
     },
+    {
+      code: dedent`
+        test('that we error on this', () => {
+          const promise = something().then(value => {
+            expect(value).toBe('red');
+          });
+
+          log(promise);
+        });
+      `,
+      errors: [
+        {
+          messageId: 'expectInFloatingPromise',
+          line: 2,
+          column: 9,
+        },
+      ],
+    },
   ],
 });
diff --git a/src/rules/valid-expect-in-promise.ts b/src/rules/valid-expect-in-promise.ts
index 33b43863b..02d933068 100644
--- a/src/rules/valid-expect-in-promise.ts
+++ b/src/rules/valid-expect-in-promise.ts
@@ -141,6 +141,40 @@ const isPromiseMethodThatUsesValue = (
   return isIdentifier(node.argument, name);
 };
 
+/**
+ * Attempts to determine if the runtime value represented by the given `identifier`
+ * is `await`ed as an argument along the given call expression
+ */
+const isValueAwaitedInArguments = (
+  name: string,
+  call: TSESTree.CallExpression,
+): boolean => {
+  let node: TSESTree.Node = call;
+
+  while (node) {
+    if (node.type === AST_NODE_TYPES.CallExpression) {
+      for (const argument of node.arguments) {
+        if (
+          argument.type === AST_NODE_TYPES.AwaitExpression &&
+          isIdentifier(argument.argument, name)
+        ) {
+          return true;
+        }
+      }
+
+      node = node.callee;
+    }
+
+    if (node.type !== AST_NODE_TYPES.MemberExpression) {
+      break;
+    }
+
+    node = node.object;
+  }
+
+  return false;
+};
+
 /**
  * Attempts to determine if the runtime value represented by the given `identifier`
  * is `await`ed or `return`ed within the given `body` of statements
@@ -163,6 +197,14 @@ const isValueAwaitedOrReturned = (
     }
 
     if (node.type === AST_NODE_TYPES.ExpressionStatement) {
+      // it's possible that we're awaiting the value as an argument
+      if (
+        node.expression.type === AST_NODE_TYPES.CallExpression &&
+        isValueAwaitedInArguments(name, node.expression)
+      ) {
+        return true;
+      }
+
       if (
         node.expression.type === AST_NODE_TYPES.AwaitExpression &&
         isPromiseMethodThatUsesValue(node.expression, identifier)