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 if the callback throws any errors, execution is immediately interrupted and an error message is received. If the callback returns falsy value, the next check will continue until truthy value 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,79 @@ export function waitFor<T>(callback: WaitForCallback<T>, options: number | WaitF
     intervalId = setInterval(checkCallback, interval)
   })
 }
+
+export type WaitUntilCallback<T> = () => T | Promise<T>
+
+export interface WaitUntilOptions extends Pick<WaitForOptions, 'interval' | 'timeout'> {}
+
+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) => {
+    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: T) => {
+      if (!result)
+        return
+
+      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<T>
+          promiseStatus = 'pending'
+          thenable.then(
+            (resolvedValue) => {
+              promiseStatus = 'resolved'
+              onResolve(resolvedValue)
+            },
+            (rejectedValue) => {
+              promiseStatus = 'rejected'
+              onReject(rejectedValue)
+            },
+          )
+        }
+        else {
+          return onResolve(result as T)
+        }
+      }
+      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,87 @@ 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('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()
+  })
+})