Improve waitFor() types

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,

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);
+    });
   });
 }