From 6c28fc00c34c76bc9d6d55779aec41908a7c042c Mon Sep 17 00:00:00 2001
From: Vladimir Sheremet <sleuths.slews0s@icloud.com>
Date: Wed, 22 Nov 2023 15:44:23 +0100
Subject: [PATCH] fix(vitest): improve vi.waitUntil type to excude falsy types

---
 packages/vitest/src/integrations/wait.ts |  6 ++++--
 test/typescript/test-d/vi.test-d.ts      | 16 ++++++++++++++++
 2 files changed, 20 insertions(+), 2 deletions(-)
 create mode 100644 test/typescript/test-d/vi.test-d.ts

diff --git a/packages/vitest/src/integrations/wait.ts b/packages/vitest/src/integrations/wait.ts
index 6f5879f055c6..85a5b8cea147 100644
--- a/packages/vitest/src/integrations/wait.ts
+++ b/packages/vitest/src/integrations/wait.ts
@@ -100,12 +100,14 @@ export type WaitUntilCallback<T> = () => T | Promise<T>
 
 export interface WaitUntilOptions extends Pick<WaitForOptions, 'interval' | 'timeout'> {}
 
+type Truthy<T> = T extends false | '' | 0 | null | undefined ? never : T
+
 export function waitUntil<T>(callback: WaitUntilCallback<T>, options: number | WaitUntilOptions = {}) {
   const { setTimeout, setInterval, clearTimeout, clearInterval } = getSafeTimers()
   const { interval = 50, timeout = 1000 } = typeof options === 'number' ? { timeout: options } : options
   const STACK_TRACE_ERROR = new Error('STACK_TRACE_ERROR')
 
-  return new Promise<T>((resolve, reject) => {
+  return new Promise<Truthy<T>>((resolve, reject) => {
     let promiseStatus: 'idle' | 'pending' | 'resolved' | 'rejected' = 'idle'
     let timeoutId: ReturnType<typeof setTimeout>
     let intervalId: ReturnType<typeof setInterval>
@@ -125,7 +127,7 @@ export function waitUntil<T>(callback: WaitUntilCallback<T>, options: number | W
       if (intervalId)
         clearInterval(intervalId)
 
-      resolve(result)
+      resolve(result as Truthy<T>)
       return true
     }
 
diff --git a/test/typescript/test-d/vi.test-d.ts b/test/typescript/test-d/vi.test-d.ts
new file mode 100644
index 000000000000..4cc99b6ac386
--- /dev/null
+++ b/test/typescript/test-d/vi.test-d.ts
@@ -0,0 +1,16 @@
+import { expectTypeOf, test, vi } from 'vitest'
+
+test('vi.waitUntil correctly resolves return type', () => {
+  expectTypeOf(vi.waitUntil(() => 'string')).resolves.toEqualTypeOf<string>()
+  expectTypeOf(vi.waitUntil(() => 1)).resolves.toEqualTypeOf<number>()
+
+  expectTypeOf(vi.waitUntil(() => false as const)).resolves.toEqualTypeOf<never>()
+  expectTypeOf(vi.waitUntil(() => '' as const)).resolves.toEqualTypeOf<never>()
+  expectTypeOf(vi.waitUntil(() => 0 as const)).resolves.toEqualTypeOf<never>()
+
+  expectTypeOf(vi.waitUntil(() => '' as '' | number)).resolves.toEqualTypeOf<number>()
+  expectTypeOf(vi.waitUntil(() => null as null | number)).resolves.toEqualTypeOf<number>()
+  expectTypeOf(vi.waitUntil(() => undefined as undefined | number)).resolves.toEqualTypeOf<number>()
+  expectTypeOf(vi.waitUntil(() => false as false | number)).resolves.toEqualTypeOf<number>()
+  expectTypeOf(vi.waitUntil(() => 0 as 0 | string)).resolves.toEqualTypeOf<string>()
+})