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

Author: macarie

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

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

@@ -69,12 +69,36 @@ export function createExpectPoll(expect: ExpectStatic): ExpectStatic['poll'] {
             let timeoutId: any
             let lastError: any
             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)
-                resolve(await assertionFunction.call(assertion, ...args))
+                try {
+                  resolve(await assertionFunction.call(assertion, ...args))
+                }
+                catch (err) {
+                  if (chai.util.flag(assertion, '_poll.assert_once')) {
+                    clearTimeout(intervalId)
+                    clearTimeout(timeoutId)
+
+                    rejectWithCause(err)
+                  }
+                  else {
+                    throw err
+                  }
+                }
                 clearTimeout(intervalId)
                 clearTimeout(timeoutId)
               }
@@ -88,17 +112,6 @@ export function createExpectPoll(expect: ExpectStatic): ExpectStatic['poll'] {
             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.')
-                }
-                reject(
-                  copyStackTrace(
-                    error,
-                    STACK_TRACE_ERROR,
-                  ),
-                )
-              }
               check()
                 .then(() => rejectWithCause(lastError))
                 .catch(e => rejectWithCause(e))