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, label?: string): Function; -export default function waitFor( - fn: CoroutineFunction, +export default function waitFor, PromiseReturn>( + fn: AsyncFunction, + label?: string +): AsyncFunction; +export default function waitFor, T>( + fn: CoroutineFunction, label?: string -): CoroutineFunction; +): CoroutineFunction; 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); + }); }); }