fix(browser): run toMatchScreenshot only once when used with expect.element (#9132)

macarie · web-flow · commit 0d2e7e3e26c2 · 2025-12-02T12:57:25.000+01:00
diff --git a/packages/browser/src/client/tester/expect-element.ts b/packages/browser/src/client/tester/expect-element.ts
@@ -29,6 +29,11 @@ function element<T extends HTMLElement | SVGElement | null | Locator>(elementOrL
       return elementOrLocator.elements() as unknown as HTMLElement
     }
 
+    if (name === 'toMatchScreenshot' && !chai.util.flag(this, '_poll.assert_once')) {
+      // `toMatchScreenshot` should only run once after the element resolves
+      chai.util.flag(this, '_poll.assert_once', true)
+    }
+
     // element selector uses prettyDOM under the hood, which is an expensive call
     // that should not be called on each failed locator attempt to avoid memory leak:
     // https://github.com/vitest-dev/vitest/issues/7139
diff --git a/packages/utils/src/timers.ts b/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))
+}
diff --git a/packages/vitest/src/integrations/chai/poll.ts b/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,47 +83,49 @@ 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 check = async () => {
-              try {
-                chai.util.flag(assertion, '_name', key)
-                const obj = await fn()
-                chai.util.flag(assertion, 'object', obj)
-                resolve(await assertionFunction.call(assertion, ...args))
-                clearTimeout(intervalId)
-                clearTimeout(timeoutId)
-              }
-              catch (err) {
-                lastError = err
-                if (!chai.util.flag(assertion, '_isLastPollAttempt')) {
-                  intervalId = setTimeout(check, interval)
+
+            let executionPhase: 'fn' | 'assertion' = 'fn'
+            let hasTimedOut = false
+
+            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)
                 }
-              }
-            }
-            timeoutId = setTimeout(() => {
-              clearTimeout(intervalId)
-              chai.util.flag(assertion, '_isLastPollAttempt', true)
-              const rejectWithCause = (error: any) => {
-                if (error.cause == null) {
-                  error.cause = new Error('Matcher did not succeed in time.')
+
+                try {
+                  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 (isLastAttempt || (executionPhase === 'assertion' && chai.util.flag(assertion, '_poll.assert_once'))) {
+                    throwWithCause(err, STACK_TRACE_ERROR)
+                  }
+
+                  await delay(interval, setTimeout)
                 }
-                reject(
-                  copyStackTrace(
-                    error,
-                    STACK_TRACE_ERROR,
-                  ),
-                )
               }
-              check()
-                .then(() => rejectWithCause(lastError))
-                .catch(e => rejectWithCause(e))
-            }, timeout)
-            check()
-          })
+            }
+            finally {
+              clearTimeout(timerId)
+            }
+          }
           let awaited = false
           test.onFinished ??= []
           test.onFinished.push(() => {
diff --git a/test/browser/fixtures/expect-dom/toMatchScreenshot.test.ts b/test/browser/fixtures/expect-dom/toMatchScreenshot.test.ts
@@ -421,4 +421,114 @@ describe('.toMatchScreenshot', () => {
 
     expect(errorMessage).matches(/^Could not capture a stable screenshot within 100ms\.$/m)
   })
+
+  // Only run this test if snapshots aren't being updated
+  test.runIf(server.config.snapshotOptions.updateSnapshot !== 'all')(
+    'runs only once after resolving the element',
+    async ({ onTestFinished }) => {
+      const filename = `toMatchScreenshot-runs-only-once-after-resolving-the-element-1`
+      const path = join(
+        '__screenshots__',
+        'toMatchScreenshot.test.ts',
+      )
+
+      const screenshotPath = join(
+        path,
+        `${filename}-${server.browser}-${server.platform}.png`
+      )
+
+      // Create baseline screenshot with original colors
+      renderTestCase([
+        'oklch(39.6% 0.141 25.723)',
+        'oklch(40.5% 0.101 131.063)',
+        'oklch(37.9% 0.146 265.522)',
+      ])
+      const locator = page.getByTestId(dataTestId)
+
+      await locator.screenshot({
+        save: true,
+        path: screenshotPath,
+      })
+
+      onTestFinished(async () => {
+        await server.commands.removeFile(screenshotPath)
+      })
+
+      // Remove element, then re-render with inverted colors after a delay
+      document.body.innerHTML = ''
+
+      const renderDelay = 500
+      setTimeout(() => {
+        renderTestCase([
+          'oklch(37.9% 0.146 265.522)',
+          'oklch(40.5% 0.101 131.063)',
+          'oklch(39.6% 0.141 25.723)',
+        ])
+      }, renderDelay)
+
+      const start = performance.now()
+
+      // Expected behavior:
+      // 1. `expect.element()` polls until element exists (~500ms)
+      // 2. `toMatchScreenshot()` runs ONCE and fails (colors don't match baseline)
+      //
+      // If `toMatchScreenshot()` polled internally, it would retry for 30s.
+      // By checking the elapsed time we verify it only ran once.
+
+      let errorMessage: string
+
+      try {
+        await expect.element(locator).toMatchScreenshot()
+      } catch (error) {
+        errorMessage = error.message
+      }
+
+      expect(typeof errorMessage).toBe('string')
+
+      const [referencePath, actualPath, diffPath] = extractToMatchScreenshotPaths(
+        errorMessage,
+        filename,
+      )
+
+      expect(typeof referencePath).toBe('string')
+      expect(typeof actualPath).toBe('string')
+      expect(typeof diffPath).toBe('string')
+
+      expect(referencePath).toMatch(new RegExp(`${screenshotPath}$`))
+
+      onTestFinished(async () => {
+        await Promise.all([
+          server.commands.removeFile(actualPath),
+          server.commands.removeFile(diffPath),
+        ])
+      })
+
+      expect(
+        errorMessage
+          .replace(/(?:\d+)(.*?)(?:0\.\d{2})/, '<pixels>$1<ratio>')
+          .replace(referencePath, '<reference>')
+          .replace(actualPath, '<actual>')
+          .replace(diffPath, '<diff>')
+      ).toMatchInlineSnapshot(`
+        expect(element).toMatchScreenshot()
+
+        Screenshot does not match the stored reference.
+        <pixels> pixels (ratio <ratio>) differ.
+
+        Reference screenshot:
+          <reference>
+
+        Actual screenshot:
+          <actual>
+
+        Diff image:
+          <diff>
+      `)
+
+      const elapsed = performance.now() - start
+
+      // Elapsed time should be lower than the default `poll`/`element` timeout
+      expect(elapsed).toBeLessThan(30_000)
+    },
+  )
 })
