feat(vitest): support vi.waitUntil method

docs/api/vi.md

@@ -777,3 +777,10 @@ test('Server started successfully', async () => {
 ```
 
 If `vi.useFakeTimers` is used, `vi.waitFor` automatically calls `vi.advanceTimersByTime(interval)` in every check callback.
+
+### vi.waitUntil
+
+- **Type:** `function waitUntil(callback: WaitUntilCallback, options?: number | WaitUntilOptions): Promise`
+- **Version**: Since Vitest 0.34.5
+
+This is similar to `vi.waitFor`, but only allows `boolean` values to be returned. If the callback returns `false`, the next check will continue until `true` is returned. This is useful when you need to wait for something to exist before taking the next step.

packages/vitest/src/integrations/vi.ts

@@ -9,7 +9,7 @@ import { resetModules, waitForImportsToResolve } from '../utils/modules'
 import { FakeTimers } from './mock/timers'
 import type { EnhancedSpy, MaybeMocked, MaybeMockedDeep, MaybePartiallyMocked, MaybePartiallyMockedDeep } from './spy'
 import { fn, isMockFunction, spies, spyOn } from './spy'
-import { waitFor } from './wait'
+import { waitFor, waitUntil } from './wait'
 
 interface VitestUtils {
   isFakeTimers(): boolean
@@ -33,6 +33,7 @@ interface VitestUtils {
   spyOn: typeof spyOn
   fn: typeof fn
   waitFor: typeof waitFor
+  waitUntil: typeof waitUntil
 
   /**
    * Run the factory before imports are evaluated. You can return a value from the factory
@@ -300,6 +301,7 @@ function createVitest(): VitestUtils {
     spyOn,
     fn,
     waitFor,
+    waitUntil,
     hoisted<T>(factory: () => T): T {
       assertTypes(factory, '"vi.hoisted" factory', ['function'])
       return factory()

packages/vitest/src/integrations/wait.ts

@@ -95,3 +95,82 @@ export function waitFor<T>(callback: WaitForCallback<T>, options: number | WaitF
     intervalId = setInterval(checkCallback, interval)
   })
 }
+
+export type WaitUntilCallback = () => boolean | Promise<boolean>
+
+export interface WaitUntilOptions extends Pick<WaitForOptions, 'interval' | 'timeout'> {}
+
+export function waitUntil(callback: WaitUntilCallback, 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<boolean>((resolve, reject) => {
+    let promiseStatus: 'idle' | 'pending' | 'resolved' | 'rejected' = 'idle'
+    let timeoutId: ReturnType<typeof setTimeout>
+    let intervalId: ReturnType<typeof setInterval>
+
+    const onReject = (error?: Error) => {
+      if (!error)
+        error = copyStackTrace(new Error('Timed out in waitUntil!'), STACK_TRACE_ERROR)
+      reject(error)
+    }
+
+    const onResolve = (result: boolean) => {
+      if (result === false)
+        return
+
+      if (typeof result !== 'boolean')
+        return onReject(new Error(`waitUntil callback must return a boolean, or a promise that resolves to a boolean, but got ${typeof result}`))
+
+      if (timeoutId)
+        clearTimeout(timeoutId)
+      if (intervalId)
+        clearInterval(intervalId)
+
+      resolve(result)
+      return true
+    }
+
+    const checkCallback = () => {
+      if (vi.isFakeTimers())
+        vi.advanceTimersByTime(interval)
+
+      if (promiseStatus === 'pending')
+        return
+      try {
+        const result = callback()
+        if (
+          result !== null
+          && typeof result === 'object'
+          && typeof (result as any).then === 'function'
+        ) {
+          const thenable = result as PromiseLike<boolean>
+          promiseStatus = 'pending'
+          thenable.then(
+            (resolvedValue) => {
+              promiseStatus = 'resolved'
+              onResolve(resolvedValue)
+            },
+            (rejectedValue) => {
+              promiseStatus = 'rejected'
+              onReject(rejectedValue)
+            },
+          )
+        }
+        else {
+          return onResolve(result as boolean)
+        }
+      }
+      catch (error) {
+        onReject(error as Error)
+      }
+    }
+
+    if (checkCallback() === true)
+      return
+
+    timeoutId = setTimeout(onReject, timeout)
+    intervalId = setInterval(checkCallback, interval)
+  })
+}

test/core/test/wait.test.ts

@@ -102,3 +102,97 @@ describe('waitFor', () => {
     vi.useRealTimers()
   })
 })
+
+describe('waitUntil', () => {
+  describe('options', () => {
+    test('timeout', async () => {
+      expect(async () => {
+        await vi.waitUntil(() => {
+          return new Promise((resolve) => {
+            setTimeout(() => {
+              resolve(true)
+            }, 100)
+          })
+        }, 50)
+      }).rejects.toThrow('Timed out in waitUntil!')
+    })
+
+    test('interval', async () => {
+      const callback = vi.fn(() => {
+        return false
+      })
+
+      await expect(
+        vi.waitUntil(callback, {
+          timeout: 60,
+          interval: 30,
+        }),
+      ).rejects.toThrowErrorMatchingInlineSnapshot('"Timed out in waitUntil!"')
+
+      expect(callback).toHaveBeenCalledTimes(2)
+    })
+  })
+
+  test('basic', async () => {
+    let result = true
+    await vi.waitUntil(() => {
+      result = !result
+      return result
+    })
+    expect(result).toBe(true)
+  })
+
+  test('async function', async () => {
+    let finished = false
+    setTimeout(() => {
+      finished = true
+    }, 50)
+    await vi.waitUntil(async () => {
+      return Promise.resolve(finished)
+    })
+  })
+
+  test('return non-boolean', async () => {
+    const fn = vi.fn(() => {
+      return 'string'
+    })
+    await expect(vi.waitUntil(fn as unknown as () => boolean, {
+      timeout: 100,
+      interval: 10,
+    })).rejects.toThrowErrorMatchingInlineSnapshot('"waitUntil callback must return a boolean, or a promise that resolves to a boolean, but got string"')
+  })
+
+  test('stacktrace correctly when callback throw error', async () => {
+    const check = () => {
+      const _a = 1
+      // @ts-expect-error test
+      _a += 1
+      return true
+    }
+    try {
+      await vi.waitUntil(check, 20)
+    }
+    catch (error) {
+      expect((error as Error).message).toMatchInlineSnapshot('"Assignment to constant variable."')
+      expect.soft((error as Error).stack).toMatch(/at check/)
+    }
+  })
+
+  test('fakeTimer works', async () => {
+    vi.useFakeTimers()
+
+    setTimeout(() => {
+      vi.advanceTimersByTime(200)
+    }, 50)
+
+    await vi.waitUntil(() => {
+      return new Promise<boolean>((resolve) => {
+        setTimeout(() => {
+          resolve(true)
+        }, 150)
+      })
+    }, 200)
+
+    vi.useRealTimers()
+  })
+})