refactor(vitest): simplify poll assertion control flow

Author: macarie

packages/utils/src/timers.ts

@@ -72,3 +72,21 @@ export function setSafeTimers(): void {
 
   (globalThis as any)[SAFE_TIMERS_SYMBOL] = timers
 }
+
+/**
+ * Returns a promise that resolves after the specified duration.
+ *
+ * @param timeout - Delay in milliseconds
+ * @param scheduler - Timer function to use, defaults to `setTimeout`. Useful for mocked timers.
+ *
+ * @example
+ * await delay(100)
+ *
+ * @example
+ * // With mocked timers
+ * const { setTimeout } = getSafeTimers()
+ * await delay(100, setTimeout)
+ */
+export function delay(timeout: number, scheduler: typeof setTimeout = setTimeout): Promise<void> {
+  return new Promise(resolve => scheduler(resolve, timeout))
+}

packages/vitest/src/integrations/chai/poll.ts

@@ -1,7 +1,7 @@
 import type { Assertion, ExpectStatic } from '@vitest/expect'
 import type { Test } from '@vitest/runner'
 import { chai } from '@vitest/expect'
-import { getSafeTimers } from '@vitest/utils/timers'
+import { delay, getSafeTimers } from '@vitest/utils/timers'
 import { getWorkerState } from '../../runtime/utils'
 
 // these matchers are not supported because they don't make sense with poll
@@ -26,6 +26,25 @@ const unsupported = [
   // resolves
 ]
 
+/**
+ * Attaches a `cause` property to the error if missing, copies the stack trace from the source, and throws.
+ *
+ * @param error - The error to throw
+ * @param source - Error to copy the stack trace from
+ *
+ * @throws Always throws the provided error with an amended stack trace
+ */
+function throwWithCause(error: any, source: Error) {
+  if (error.cause == null) {
+    error.cause = new Error('Matcher did not succeed in time.')
+  }
+
+  throw copyStackTrace(
+    error,
+    source,
+  )
+}
+
 export function createExpectPoll(expect: ExpectStatic): ExpectStatic['poll'] {
   return function poll(fn, options = {}) {
     const state = getWorkerState()
@@ -64,60 +83,52 @@ export function createExpectPoll(expect: ExpectStatic): ExpectStatic['poll'] {
 
         return function (this: any, ...args: any[]) {
           const STACK_TRACE_ERROR = new Error('STACK_TRACE_ERROR')
-          const promise = () => new Promise<void>((resolve, reject) => {
-            let intervalId: any
-            let timeoutId: any
-            let lastError: any
+          const promise = async () => {
             const { setTimeout, clearTimeout } = getSafeTimers()
-            const rejectWithCause = (error: any) => {
-              if (error.cause == null) {
-                error.cause = new Error('Matcher did not succeed in time.')
-              }
-              reject(
-                copyStackTrace(
-                  error,
-                  STACK_TRACE_ERROR,
-                ),
-              )
-            }
-            const check = async () => {
-              try {
-                chai.util.flag(assertion, '_name', key)
-                const obj = await fn()
-                chai.util.flag(assertion, 'object', obj)
+
+            let executionPhase: 'fn' | 'assertion' = 'fn'
+            let hasTimedOut = false
+            let lastError: any
+
+            const timerId = setTimeout(() => {
+              hasTimedOut = true
+            }, timeout)
+
+            chai.util.flag(assertion, '_name', key)
+
+            try {
+              while (true) {
+                const isLastAttempt = hasTimedOut
+
+                if (isLastAttempt) {
+                  chai.util.flag(assertion, '_isLastPollAttempt', true)
+                }
+
                 try {
-                  resolve(await assertionFunction.call(assertion, ...args))
+                  executionPhase = 'fn'
+                  const obj = await fn()
+                  chai.util.flag(assertion, 'object', obj)
+
+                  executionPhase = 'assertion'
+                  const output = await assertionFunction.call(assertion, ...args)
+
+                  return output
                 }
                 catch (err) {
-                  if (chai.util.flag(assertion, '_poll.assert_once')) {
-                    clearTimeout(intervalId)
-                    clearTimeout(timeoutId)
+                  lastError = err
 
-                    rejectWithCause(err)
-                  }
-                  else {
-                    throw err
+                  if (isLastAttempt || (executionPhase === 'assertion' && chai.util.flag(assertion, '_poll.assert_once'))) {
+                    throwWithCause(lastError, STACK_TRACE_ERROR)
                   }
-                }
-                clearTimeout(intervalId)
-                clearTimeout(timeoutId)
-              }
-              catch (err) {
-                lastError = err
-                if (!chai.util.flag(assertion, '_isLastPollAttempt')) {
-                  intervalId = setTimeout(check, interval)
+
+                  await delay(interval, setTimeout)
                 }
               }
             }
-            timeoutId = setTimeout(() => {
-              clearTimeout(intervalId)
-              chai.util.flag(assertion, '_isLastPollAttempt', true)
-              check()
-                .then(() => rejectWithCause(lastError))
-                .catch(e => rejectWithCause(e))
-            }, timeout)
-            check()
-          })
+            finally {
+              clearTimeout(timerId)
+            }
+          }
           let awaited = false
           test.onFinished ??= []
           test.onFinished.push(() => {