feat(vitest): support vi.waitFor method (#4113)

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

Author: Dunqing

docs/api/vi.md

@@ -254,10 +254,10 @@ import { vi } from 'vitest'
   ```ts
   // increment.test.js
   import { vi } from 'vitest'
-  
+
   // axios is a default export from `__mocks__/axios.js`
   import axios from 'axios'
-  
+
   // increment is a named export from `src/__mocks__/increment.js`
   import { increment } from '../increment.js'
 
@@ -371,7 +371,7 @@ test('importing the next module imports mocked one', async () => {
 
   ```ts
   import { vi } from 'vitest'
-  
+
   import { data } from './data.js' // Will not get reevaluated beforeEach test
 
   beforeEach(() => {
@@ -706,8 +706,74 @@ unmockedIncrement(30) === 31
 
   The implementation is based internally on [`@sinonjs/fake-timers`](https://github.com/sinonjs/fake-timers).
 
+## vi.isFakeTimers
+
+- **Type:** `() => boolean`
+- **Version:** Since Vitest 0.34.5
+
+  Returns `true` if fake timers are enabled.
+
 ## vi.useRealTimers
 
 - **Type:** `() => Vitest`
 
   When timers are run out, you may call this method to return mocked timers to its original implementations. All timers that were run before will not be restored.
+
+### vi.waitFor
+
+- **Type:** `function waitFor<T>(callback: WaitForCallback<T>, options?: number | WaitForOptions): Promise<T>`
+- **Version**: Since Vitest 0.34.5
+
+Wait for the callback to execute successfully. If the callback throws an error or returns a rejected promise it will continue to wait until it succeeds or times out.
+
+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'
+
+test('Server started successfully', async () => {
+  let server = false
+
+  setTimeout(() => {
+    server = true
+  }, 100)
+
+  function checkServerStart() {
+    if (!server)
+      throw new Error('Server not started')
+
+    console.log('Server started')
+  }
+
+  const res = await vi.waitFor(checkServerStart, {
+    timeout: 500, // default is 1000
+    interval: 20, // default is 50
+  })
+  expect(server).toBe(true)
+})
+```
+
+It also works for asynchronous callbacks
+
+```ts
+import { test, vi } from 'vitest'
+
+test('Server started successfully', async () => {
+  async function startServer() {
+    return new Promise((resolve) => {
+      setTimeout(() => {
+        server = true
+        resolve('Server started')
+      }, 100)
+    })
+  }
+
+  const server = await vi.waitFor(startServer, {
+    timeout: 500, // default is 1000
+    interval: 20, // default is 50
+  })
+  expect(server).toBe('Server started')
+})
+```
+
+If `vi.useFakeTimers` is used, `vi.waitFor` automatically calls `vi.advanceTimersByTime(interval)` in every check callback.

packages/vitest/src/integrations/mock/timers.ts

@@ -175,6 +175,10 @@ export class FakeTimers {
     this._userConfig = config
   }
 
+  isFakeTimers() {
+    return this._fakingTime
+  }
+
   private _checkFakeTimers() {
     if (!this._fakingTime) {
       throw new Error(

packages/vitest/src/integrations/vi.ts

@@ -9,8 +9,10 @@ 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'
 
 interface VitestUtils {
+  isFakeTimers(): boolean
   useFakeTimers(config?: FakeTimerInstallOpts): this
   useRealTimers(): this
   runOnlyPendingTimers(): this
@@ -30,6 +32,7 @@ interface VitestUtils {
 
   spyOn: typeof spyOn
   fn: typeof fn
+  waitFor: typeof waitFor
 
   /**
    * Run the factory before imports are evaluated. You can return a value from the factory
@@ -213,6 +216,10 @@ function createVitest(): VitestUtils {
       return utils
     },
 
+    isFakeTimers() {
+      return _timers.isFakeTimers()
+    },
+
     useRealTimers() {
       _timers.useRealTimers()
       _mockedDate = null
@@ -292,7 +299,7 @@ function createVitest(): VitestUtils {
 
     spyOn,
     fn,
-
+    waitFor,
     hoisted<T>(factory: () => T): T {
       assertTypes(factory, '"vi.hoisted" factory', ['function'])
       return factory()

packages/vitest/src/integrations/wait.ts

@@ -0,0 +1,97 @@
+import { getSafeTimers } from '@vitest/utils'
+import { vi } from './vi'
+
+// The waitFor function was inspired by https://github.com/testing-library/web-testing-library/pull/2
+
+export type WaitForCallback<T> = () => T | Promise<T>
+
+export interface WaitForOptions {
+  /**
+   * @description Time in ms between each check callback
+   * @default 50ms
+   */
+  interval?: number
+  /**
+   * @description Time in ms after which the throw a timeout error
+   * @default 1000ms
+   */
+  timeout?: number
+}
+
+function copyStackTrace(target: Error, source: Error) {
+  if (source.stack !== undefined)
+    target.stack = source.stack.replace(source.message, target.message)
+  return target
+}
+
+export function waitFor<T>(callback: WaitForCallback<T>, options: number | WaitForOptions = {}) {
+  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 lastError: unknown
+    let promiseStatus: 'idle' | 'pending' | 'resolved' | 'rejected' = 'idle'
+    let timeoutId: ReturnType<typeof setTimeout>
+    let intervalId: ReturnType<typeof setInterval>
+
+    const onResolve = (result: T) => {
+      if (timeoutId)
+        clearTimeout(timeoutId)
+      if (intervalId)
+        clearInterval(intervalId)
+
+      resolve(result)
+    }
+
+    const handleTimeout = () => {
+      let error = lastError
+      if (!error)
+        error = copyStackTrace(new Error('Timed out in waitFor!'), STACK_TRACE_ERROR)
+
+      reject(error)
+    }
+
+    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'
+              lastError = rejectedValue
+            },
+          )
+        }
+        else {
+          onResolve(result as T)
+          return true
+        }
+      }
+      catch (error) {
+        lastError = error
+      }
+    }
+
+    if (checkCallback() === true)
+      return
+
+    timeoutId = setTimeout(handleTimeout, timeout)
+    intervalId = setInterval(checkCallback, interval)
+  })
+}

test/core/test/wait.test.ts

@@ -0,0 +1,104 @@
+import { describe, expect, test, vi } from 'vitest'
+
+describe('waitFor', () => {
+  describe('options', () => {
+    test('timeout', async () => {
+      expect(async () => {
+        await vi.waitFor(() => {
+          return new Promise((resolve) => {
+            setTimeout(() => {
+              resolve(true)
+            }, 100)
+          })
+        }, 50)
+      }).rejects.toThrow('Timed out in waitFor!')
+    })
+
+    test('interval', async () => {
+      const callback = vi.fn(() => {
+        throw new Error('interval error')
+      })
+
+      await expect(
+        vi.waitFor(callback, {
+          timeout: 60,
+          interval: 30,
+        }),
+      ).rejects.toThrowErrorMatchingInlineSnapshot('"interval error"')
+
+      expect(callback).toHaveBeenCalledTimes(2)
+    })
+  })
+
+  test('basic', async () => {
+    let throwError = false
+    await vi.waitFor(() => {
+      if (!throwError) {
+        throwError = true
+        throw new Error('basic error')
+      }
+    })
+    expect(throwError).toBe(true)
+  })
+
+  test('async function', async () => {
+    let finished = false
+    setTimeout(() => {
+      finished = true
+    }, 50)
+    await vi.waitFor(async () => {
+      if (finished)
+        return Promise.resolve(true)
+      else
+        return Promise.reject(new Error('async function error'))
+    })
+  })
+
+  test('stacktrace correctly', async () => {
+    const check = () => {
+      const _a = 1
+      // @ts-expect-error test
+      _a += 1
+    }
+    try {
+      await vi.waitFor(check, 100)
+    }
+    catch (error) {
+      expect((error as Error).message).toMatchInlineSnapshot('"Assignment to constant variable."')
+      expect.soft((error as Error).stack).toMatch(/at check/)
+    }
+  })
+
+  test('stacktrace point to waitFor', async () => {
+    const check = async () => {
+      return new Promise((resolve) => {
+        setTimeout(resolve, 60)
+      })
+    }
+    try {
+      await vi.waitFor(check, 50)
+    }
+    catch (error) {
+      expect(error).toMatchInlineSnapshot('[Error: Timed out in waitFor!]')
+      expect((error as Error).stack?.split('\n')[1]).toMatch(/waitFor\s*\(.*\)?/)
+    }
+  })
+
+  test('fakeTimer works', async () => {
+    vi.useFakeTimers()
+
+    setTimeout(() => {
+      vi.advanceTimersByTime(200)
+    }, 50)
+
+    await vi.waitFor(() => {
+      return new Promise<void>((resolve) => {
+        setTimeout(() => {
+          resolve()
+        }, 150)
+      })
+    }, 200)
+
+    vi.useRealTimers()
+  })
+})