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

tsirlucas · web-flow · commit 847d32211929 · 2025-01-03T16:08:35.000+01:00
diff --git a/packages/browser/src/client/tester/expect-element.ts b/packages/browser/src/client/tester/expect-element.ts
@@ -18,11 +18,26 @@ export async function setupExpectDom() {
 
       const isNot = chai.util.flag(this, 'negate') as boolean
       const name = chai.util.flag(this, '_name') as string
+      // 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
+      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)
   }
 }
diff --git a/packages/vitest/src/integrations/chai/poll.ts b/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,28 @@ 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(() => {
+              clearTimeout(intervalId)
+              chai.util.flag(assertion, '_isLastPollAttempt', true)
+              const rejectWithCause = (cause: any) => {
+                reject(
+                  copyStackTrace(
+                    new Error(`Matcher did not succeed in ${timeout}ms`, {
+                      cause,
+                    }),
+                    STACK_TRACE_ERROR,
+                  ),
+                )
+              }
+              check()
+                .then(() => rejectWithCause(lastError))
+                .catch(e => rejectWithCause(e))
+            }, timeout)
             check()
           })
           let awaited = false
diff --git a/test/browser/specs/runner.test.ts b/test/browser/specs/runner.test.ts
@@ -28,8 +28,11 @@ describe('running browser tests', async () => {
       console.error(stderr)
     })
 
-    expect(browserResultJson.testResults).toHaveLength(19 * instances.length)
-    expect(passedTests).toHaveLength(17 * instances.length)
+    // This should match the number of actual tests from browser.json
+    // if you added new tests, these assertion will fail and you should
+    // update the numbers
+    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')
diff --git a/test/browser/test/expect-element.test.ts b/test/browser/test/expect-element.test.ts
@@ -0,0 +1,24 @@
+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.spyOn(locator, 'element')
+  const locatorQueryMock = vi.spyOn(locator, 'query')
+
+  try {
+    await expect.element(locator, { timeout: 500, interval: 100 }).toBeInTheDocument()
+  }
+  catch {}
+
+  expect(locatorElementMock).toBeCalledTimes(1)
+  expect(locatorElementMock).toHaveBeenCalledAfter(locatorQueryMock)
+})
diff --git a/test/core/test/expect-poll.test.ts b/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,44 @@ 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)
+  })
+})
+
+test('should handle success on last attempt', async () => {
+  const fn = vi.fn(function (this: object) {
+    if (chai.util.flag(this, '_isLastPollAttempt')) {
+      return 1
+    }
+    return undefined
+  })
+  await expect.poll(fn, { interval: 100, timeout: 500 }).toBe(1)
+})
+
+test('should handle failure on last attempt', async () => {
+  const fn = vi.fn(function (this: object) {
+    if (chai.util.flag(this, '_isLastPollAttempt')) {
+      return 3
+    }
+    return 2
+  })
+  await expect(async () => {
+    await expect.poll(fn, { interval: 10, timeout: 100 }).toBe(1)
+  }).rejects.toThrowError(expect.objectContaining({
+    message: 'Matcher did not succeed in 100ms',
+    cause: expect.objectContaining({
+      // makes sure cause message reflects the last attempt value
+      message: 'expected 3 to be 1 // Object.is equality',
+    }),
+  }))
+})
