From 216138129323fe5638a78f4253450ce1f363cf50 Mon Sep 17 00:00:00 2001
From: Ben Demboski <demboskb@gmail.com>
Date: Fri, 21 Jul 2023 14:10:12 -0700
Subject: [PATCH] Improve waitFor() types

Use type arguments & inference to ensure that `waitFor()` in function form returns the same type that it was passed. This prevents it from getting in the way of passing functions as as typed arguments. For example, `ember-concurrency`'s `task()` function expects an async arrow function as its argument, so `task(async () => null)` works, but without this change, `task(waitFor(() => null))` does not because `waitFor(() => null)` is just `Function`.
---
 addon/@ember/test-waiters/wait-for.ts | 11 +++++++----
 tests/unit/wait-for-test.ts           | 27 +++++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/addon/@ember/test-waiters/wait-for.ts b/addon/@ember/test-waiters/wait-for.ts
index abdef484..a92f101f 100644
--- a/addon/@ember/test-waiters/wait-for.ts
+++ b/addon/@ember/test-waiters/wait-for.ts
@@ -75,11 +75,14 @@ type DecoratorArguments = [object, string, PropertyDescriptor, string?];
  * }
  *
  */
-export default function waitFor(fn: AsyncFunction<any[], any>, label?: string): Function;
-export default function waitFor(
-  fn: CoroutineFunction<any[], any>,
+export default function waitFor<A extends Array<any>, PromiseReturn>(
+  fn: AsyncFunction<A, PromiseReturn>,
+  label?: string
+): AsyncFunction<A, PromiseReturn>;
+export default function waitFor<A extends Array<any>, T>(
+  fn: CoroutineFunction<A, T>,
   label?: string
-): CoroutineFunction<any[], any>;
+): CoroutineFunction<A, T>;
 export default function waitFor(
   target: object,
   _key: string,
diff --git a/tests/unit/wait-for-test.ts b/tests/unit/wait-for-test.ts
index 78572246..df51b175 100644
--- a/tests/unit/wait-for-test.ts
+++ b/tests/unit/wait-for-test.ts
@@ -445,5 +445,32 @@ if (DEBUG) {
         assert.deepEqual(getPendingWaiterState().pending, 0);
       });
     });
+
+    test('types', async function (assert) {
+      assert.expect(0);
+
+      async function asyncFn(a: string, b: string) {
+        return `${a}${b}`;
+      }
+      function* genFn(a: string, b: string) {
+        yield `${a}${b}`;
+        return `${a}${b}`;
+      }
+
+      function asyncNoop(fn: typeof asyncFn) {
+        return fn;
+      }
+      function genNoop(fn: typeof genFn) {
+        return fn;
+      }
+
+      asyncNoop(waitFor(asyncFn));
+      genNoop(waitFor(genFn));
+
+      // @ts-expect-error wrong argument types
+      waitFor(asyncFn)(1, 2);
+      // @ts-expect-error wrong argument types
+      waitFor(genFn)(1, 2);
+    });
   });
 }