fix: only use locator.element on last expect.element attempt (fix #7139)

packages/browser/src/client/tester/expect-element.ts

@@ -18,11 +18,23 @@ export async function setupExpectDom() {
 
       const isNot = chai.util.flag(this, 'negate') as boolean
       const name = chai.util.flag(this, '_name') as string
+      const isLastPollAttempt = chai.util.flag(this, '_isLastPollAttempt')
       // special case for `toBeInTheDocument` matcher
       if (isNot && name === 'toBeInTheDocument') {
         return elementOrLocator.query()
       }
-      return elementOrLocator.element()
+
+      if (isLastPollAttempt) {
+        return elementOrLocator.element()
+      }
+
+      const result = elementOrLocator.query()
+
+      if (!result) {
+        throw new Error(`Cannot find element with locator: ${JSON.stringify(elementOrLocator)}`)
+      }
+
+      return result
     }, options)
   }
 }

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

@@ -66,19 +66,9 @@ export function createExpectPoll(expect: ExpectStatic): ExpectStatic['poll'] {
           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 { setTimeout, clearTimeout } = getSafeTimers()
-            const timeoutId = setTimeout(() => {
-              clearTimeout(intervalId)
-              reject(
-                copyStackTrace(
-                  new Error(`Matcher did not succeed in ${timeout}ms`, {
-                    cause: lastError,
-                  }),
-                  STACK_TRACE_ERROR,
-                ),
-              )
-            }, timeout)
             const check = async () => {
               try {
                 chai.util.flag(assertion, '_name', key)
@@ -90,9 +80,24 @@ export function createExpectPoll(expect: ExpectStatic): ExpectStatic['poll'] {
               }
               catch (err) {
                 lastError = err
-                intervalId = setTimeout(check, interval)
+                if (!chai.util.flag(assertion, '_isLastPollAttempt')) {
+                  intervalId = setTimeout(check, interval)
+                }
               }
             }
+            timeoutId = setTimeout(async () => {
+              clearTimeout(intervalId)
+              chai.util.flag(assertion, '_isLastPollAttempt', true)
+              await check()
+              reject(
+                copyStackTrace(
+                  new Error(`Matcher did not succeed in ${timeout}ms`, {
+                    cause: lastError,
+                  }),
+                  STACK_TRACE_ERROR,
+                ),
+              )
+            }, timeout)
             check()
           })
           let awaited = false

test/browser/specs/runner.test.ts

@@ -28,8 +28,8 @@ describe('running browser tests', async () => {
       console.error(stderr)
     })
 
-    expect(browserResultJson.testResults).toHaveLength(19 * instances.length)
-    expect(passedTests).toHaveLength(17 * instances.length)
+    expect(browserResultJson.testResults).toHaveLength(20 * instances.length)
+    expect(passedTests).toHaveLength(18 * instances.length)
     expect(failedTests).toHaveLength(2 * instances.length)
 
     expect(stderr).not.toContain('optimized dependencies changed')

test/browser/test/expect-element.test.ts

@@ -0,0 +1,26 @@
+import { page } from '@vitest/browser/context'
+import { expect, test, vi } from 'vitest'
+
+// 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
+test('should only use element selector on last expect.element attempt', async () => {
+  const div = document.createElement('div')
+  const spanString = '<span>test</span>'
+  div.innerHTML = spanString
+  document.body.append(div)
+
+  const locator = page.getByText('non-existent')
+  const locatorElementMock = vi.fn(locator.element)
+  locator.element = locatorElementMock
+  const locatorQueryMock = vi.fn(locator.query)
+  locator.query = locatorQueryMock
+
+  try {
+    await expect.element(locator, { timeout: 500, interval: 100 }).toBeInTheDocument()
+  }
+  catch {}
+
+  expect(locatorElementMock).toBeCalledTimes(1)
+  expect(locatorElementMock).toHaveBeenCalledAfter(locatorQueryMock)
+})

test/core/test/expect-poll.test.ts

@@ -1,4 +1,4 @@
-import { expect, test, vi } from 'vitest'
+import { chai, expect, test, vi } from 'vitest'
 
 test('simple usage', async () => {
   await expect.poll(() => false).toBe(false)
@@ -106,3 +106,16 @@ test('toBeDefined', async () => {
     }),
   }))
 })
+
+test('should set _isLastPollAttempt flag on last call', async () => {
+  const fn = vi.fn(function (this: object) {
+    return chai.util.flag(this, '_isLastPollAttempt')
+  })
+  await expect(async () => {
+    await expect.poll(fn, { interval: 100, timeout: 500 }).toBe(false)
+  }).rejects.toThrowError()
+  fn.mock.results.forEach((result, index) => {
+    const isLastCall = index === fn.mock.results.length - 1
+    expect(result.value).toBe(isLastCall ? true : undefined)
+  })
+})