feat(vitest): support vi.waitUntil method (#4129)

Co-authored-by: Vladimir Sheremet <[email protected]>

Author: Dunqing

docs/api/vi.md

@@ -729,7 +729,7 @@ Wait for the callback to execute successfully. If the callback throws an error o
 This is very useful when you need to wait for some asynchronous action to complete, for example, when you start a server and need to wait for it to start.
 
 ```ts
-import { test, vi } from 'vitest'
+import { expect, test, vi } from 'vitest'
 
 test('Server started successfully', async () => {
   let server = false
@@ -756,7 +756,7 @@ test('Server started successfully', async () => {
 It also works for asynchronous callbacks
 
 ```ts
-import { test, vi } from 'vitest'
+import { expect, test, vi } from 'vitest'
 
 test('Server started successfully', async () => {
   async function startServer() {
@@ -777,3 +777,29 @@ 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.
+
+Look at the example below. We can use `vi.waitUntil` to wait for the element to appear on the page, and then we can do something with the element.
+
+```ts
+import { expect, test, vi } from 'vitest'
+
+test('Element render correctly', async () => {
+  const element = await vi.waitUntil(
+    () => document.querySelector('.element'),
+    {
+      timeout: 500, // default is 1000
+      interval: 20, // default is 50
+    }
+  )
+
+  // do something with the element
+  expect(element.querySelector('.element-child')).toBeTruthy()
+})
+```

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