From 4df048c1131a6c484965467e96254d930b17e41c Mon Sep 17 00:00:00 2001 From: Hiroshi Ogawa Date: Fri, 8 May 2026 00:51:43 +0900 Subject: [PATCH] fix!: fail `expect.poll` when function didn't resolve in time (#10233) Co-authored-by: Codex --- docs/api/browser/assertions.md | 2 +- docs/api/expect.md | 6 +- packages/browser/matchers.d.ts | 3 +- .../src/client/tester/expect-element.ts | 31 ++------ packages/snapshot/src/client.ts | 10 ++- packages/vitest/src/integrations/chai/poll.ts | 79 +++++++++++++------ packages/vitest/src/types/global.ts | 2 +- test/browser/test/expect-element.test.ts | 24 ------ test/browser/test/findElement.test.ts | 13 +++ test/snapshots/test/domain-poll.test.ts | 50 ++++++++++++ test/unit/test/expect-poll.test.ts | 78 +++++++++--------- 11 files changed, 185 insertions(+), 113 deletions(-) delete mode 100644 test/browser/test/expect-element.test.ts diff --git a/docs/api/browser/assertions.md b/docs/api/browser/assertions.md index c6ae18b6a..c1c0cd078 100644 --- a/docs/api/browser/assertions.md +++ b/docs/api/browser/assertions.md @@ -54,7 +54,7 @@ interface ExpectPollOptions { ``` ::: tip -`expect.element` is a shorthand for `expect.poll(() => element)` and works in exactly the same way. +Like [`expect.poll`](/api/expect#poll), `expect.element` retries DOM assertions until they pass or the timeout is reached. When it receives a locator, Vitest resolves it with [`locator.findElement()`](/api/browser/locators#findelement) before running the DOM assertion. The `timeout` option applies to the whole retry operation. The `interval` option controls how often failed DOM assertions are retried, but locator resolution uses `findElement`'s own increasing retry intervals. `toHaveTextContent` and all other assertions are still available on a regular `expect` without a built-in retry-ability mechanism: diff --git a/docs/api/expect.md b/docs/api/expect.md index c5ac57ee8..d2603aa80 100644 --- a/docs/api/expect.md +++ b/docs/api/expect.md @@ -105,11 +105,13 @@ test('expect.soft test', () => { ```ts interface ExpectPoll extends ExpectStatic { - (actual: () => T, options?: { interval?: number; timeout?: number; message?: string }): Promise> + (actual: (options: { signal: AbortSignal }) => T, options?: { interval?: number; timeout?: number; message?: string }): Promise>> } ``` -`expect.poll` reruns the _assertion_ until it is succeeded. You can configure how many times Vitest should rerun the `expect.poll` callback by setting `interval` and `timeout` options. +`expect.poll` reruns the _assertion_ until it is succeeded. You can configure how often Vitest retries and how long it waits by setting `interval` and `timeout` options. The `timeout` applies to the whole polling operation, including pending callback and async matcher execution. + +The callback receives an `AbortSignal` that is aborted when the poll timeout is reached. If an error is thrown inside the `expect.poll` callback, Vitest will retry again until the timeout runs out. diff --git a/packages/browser/matchers.d.ts b/packages/browser/matchers.d.ts index 75d2d512e..ba0caae97 100644 --- a/packages/browser/matchers.d.ts +++ b/packages/browser/matchers.d.ts @@ -18,7 +18,8 @@ declare module 'vitest' { interface ExpectStatic { /** - * `expect.element(locator)` is a shorthand for `expect.poll(() => locator.element())`. + * `expect.element(locator)` retries locator resolution and DOM assertions + * using the `expect.poll` timeout options. * You can set default timeout via `expect.poll.timeout` option in the config. * @see {@link https://vitest.dev/api/expect#poll} */ diff --git a/packages/browser/src/client/tester/expect-element.ts b/packages/browser/src/client/tester/expect-element.ts index a2d93a66a..3b2474d9b 100644 --- a/packages/browser/src/client/tester/expect-element.ts +++ b/packages/browser/src/client/tester/expect-element.ts @@ -16,7 +16,9 @@ function element(elementOrL throw new Error(`Invalid element or locator: ${elementOrLocator}. Expected an instance of HTMLElement, SVGElement or Locator, received ${getType(elementOrLocator)}`) } - const expectElement = expect.poll(function element(this: object) { + const pollOptions = processTimeoutOptions(options) + const deadline = pollOptions?.timeout ? now() + pollOptions.timeout : undefined + const expectElement = expect.poll(async function element(this: object): Promise { if (elementOrLocator instanceof Element || elementOrLocator == null) { return elementOrLocator } @@ -33,28 +35,11 @@ function element(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 - const isLastPollAttempt = chai.util.flag(this, '_isLastPollAttempt') - - if (isLastPollAttempt) { - return elementOrLocator.element() - } - - const result = elementOrLocator.query() - - if (!result) { - throw new Error(`Cannot find element with locator: ${JSON.stringify(elementOrLocator)}`) - } - - return result - }, processTimeoutOptions(options)) + return elementOrLocator.findElement({ + ...pollOptions, + timeout: deadline ? Math.max(deadline - now(), 0) : undefined, + }) + }, pollOptions) chai.util.flag(expectElement, '_poll.element', true) diff --git a/packages/snapshot/src/client.ts b/packages/snapshot/src/client.ts index c79a50280..edf4502a0 100644 --- a/packages/snapshot/src/client.ts +++ b/packages/snapshot/src/client.ts @@ -58,7 +58,7 @@ interface AssertDomainOptions extends Omit { } interface AssertDomainPollOptions extends Omit { - poll: () => Promise | unknown + poll: (options: { signal: AbortSignal }) => Promise | unknown timeout?: number interval?: number } @@ -294,12 +294,16 @@ export class SnapshotClient { const reference = expectedSnapshot.data !== undefined && snapshotState.snapshotUpdateState !== 'all' ? adapter.parseExpected(expectedSnapshot.data) : undefined + const timeoutController = new AbortController() const timedOut = timeout > 0 - ? new Promise(r => setTimeout(r, timeout)) + ? new Promise(r => setTimeout(() => { + timeoutController.abort() + r() + }, timeout)) : undefined const stableResult = await getStableSnapshot({ adapter, - poll, + poll: () => poll({ signal: timeoutController.signal }), interval, timedOut, match: reference diff --git a/packages/vitest/src/integrations/chai/poll.ts b/packages/vitest/src/integrations/chai/poll.ts index a1cf62931..2e214dd44 100644 --- a/packages/vitest/src/integrations/chai/poll.ts +++ b/packages/vitest/src/integrations/chai/poll.ts @@ -133,40 +133,57 @@ export function createExpectPoll(expect: ExpectStatic): ExpectStatic['poll'] { } const { setTimeout, clearTimeout } = getSafeTimers() - - let executionPhase: 'fn' | 'assertion' = 'fn' - let hasTimedOut = false - - const timerId = setTimeout(() => { - hasTimedOut = true - }, timeout) + let timerId: ReturnType | undefined + const timeoutController = new AbortController() + const timeoutPromise = new Promise((resolve) => { + timerId = setTimeout(() => { + timeoutController.abort() + resolve() + }, timeout) + }) + let lastError: unknown try { while (true) { - const isLastAttempt = hasTimedOut - - if (isLastAttempt) { - chai.util.flag(assertion, '_isLastPollAttempt', true) - } - try { - executionPhase = 'fn' - const obj = await fn() + const fnResult = await raceWith( + Promise.resolve().then(() => fn({ signal: timeoutController.signal })), + timeoutPromise, + ) + if (!fnResult.ok) { + lastError ??= new Error(`expect.poll() function didn't resolve in time.`) + break + } + const obj = fnResult.value chai.util.flag(assertion, 'object', obj) - executionPhase = 'assertion' - const output = await assertionFunction.call(assertion, ...args) + const assertionResult = await raceWith( + Promise.resolve().then(() => assertionFunction.apply(assertion, args)), + timeoutPromise, + ) + if (!assertionResult.ok) { + lastError ??= new Error(`expect.poll() assertion didn't resolve in time.`) + break + } + const output = assertionResult.value await onSettled?.({ assertion, status: 'pass' }) return output } catch (err) { - if (isLastAttempt || (executionPhase === 'assertion' && chai.util.flag(assertion, '_poll.assert_once'))) { - await onSettled?.({ assertion, status: 'fail' }) - throwWithCause(err, STACK_TRACE_ERROR) + lastError = err + // no retry for toMatchScreenshot since + // it owns retry/stability after the first element resolution + if (key === 'toMatchScreenshot') { + break + } + const result = await raceWith( + delay(interval, setTimeout), + timeoutPromise, + ) + if (!result.ok) { + break } - - await delay(interval, setTimeout) if (vi.isFakeTimers()) { vi.advanceTimersByTime(interval) } @@ -176,6 +193,10 @@ export function createExpectPoll(expect: ExpectStatic): ExpectStatic['poll'] { finally { clearTimeout(timerId) } + if (lastError) { + await onSettled?.({ assertion, status: 'fail' }) + throwWithCause(lastError, STACK_TRACE_ERROR) + } } let awaited = false test.onFinished ??= [] @@ -221,3 +242,17 @@ function copyStackTrace(target: Error, source: Error) { } return target } + +function raceWith( + promise: Promise, + other?: Promise, +): Promise<{ ok: true; value: A } | { ok: false; value: B }> { + const left = promise.then(value => ({ ok: true as const, value })) + if (!other) { + return left + } + return Promise.race([ + left, + other.then(value => ({ ok: false as const, value })), + ]) +} diff --git a/packages/vitest/src/types/global.ts b/packages/vitest/src/types/global.ts index d22bfa5df..9ac1fa9a7 100644 --- a/packages/vitest/src/types/global.ts +++ b/packages/vitest/src/types/global.ts @@ -40,7 +40,7 @@ declare module 'vitest' { unreachable: (message?: string) => never soft: (actual: T, message?: string) => Assertion poll: ( - actual: () => T, + actual: (options: { signal: AbortSignal }) => T, options?: ExpectPollOptions, ) => PromisifyAssertion> addEqualityTesters: (testers: Array) => void diff --git a/test/browser/test/expect-element.test.ts b/test/browser/test/expect-element.test.ts deleted file mode 100644 index aba9ce83c..000000000 --- a/test/browser/test/expect-element.test.ts +++ /dev/null @@ -1,24 +0,0 @@ -import { expect, test, vi } from 'vitest' -import { page } from 'vitest/browser' - -// 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 = 'test' - 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/browser/test/findElement.test.ts b/test/browser/test/findElement.test.ts index a9e8a3f25..3a09cb761 100644 --- a/test/browser/test/findElement.test.ts +++ b/test/browser/test/findElement.test.ts @@ -108,3 +108,16 @@ function createButton() { document.body.append(button) return button } + +test('expect.element is strict', async () => { + createButton() + createButton() + await expect( + () => expect.element(page.getByRole('button'), { timeout: 50 }).toBeVisible(), + ).rejects.toThrowErrorMatchingInlineSnapshot(` + [Error: strict mode violation: getByRole('button') resolved to 2 elements: + 1) aka getByRole('button').first() + 2) aka getByRole('button').nth(1) + ] + `) +}) diff --git a/test/snapshots/test/domain-poll.test.ts b/test/snapshots/test/domain-poll.test.ts index 46ed59f5a..bfc38d6df 100644 --- a/test/snapshots/test/domain-poll.test.ts +++ b/test/snapshots/test/domain-poll.test.ts @@ -444,3 +444,53 @@ test('throwing', async () => { } `) }) + +test('signal', async () => { + const result = await runInlineTests({ + 'basic.test.ts': ` +import '../test/fixtures/domain/basic-extend' + +test('signal', async () => { + let aborted = false + await expect( + expect.poll(({ signal }) => { + signal.addEventListener('abort', () => { + aborted = true + }) + return new Promise(() => {}) + }, { timeout: 100, interval: 10 }).toMatchKvSnapshot() + ).rejects.toThrowErrorMatchingInlineSnapshot() + expect(aborted).toMatchInlineSnapshot() +}) +`, + }, { + globals: true, + update: 'all', + }) + expect(result.stderr).toMatchInlineSnapshot(`""`) + expect(result.errorTree()).toMatchInlineSnapshot(` + Object { + "basic.test.ts": Object { + "signal": "passed", + }, + } + `) + expect(result.fs.readFile('basic.test.ts')).toMatchInlineSnapshot(` + " + import '../test/fixtures/domain/basic-extend' + + test('signal', async () => { + let aborted = false + await expect( + expect.poll(({ signal }) => { + signal.addEventListener('abort', () => { + aborted = true + }) + return new Promise(() => {}) + }, { timeout: 100, interval: 10 }).toMatchKvSnapshot() + ).rejects.toThrowErrorMatchingInlineSnapshot(\`[Error: poll() did not produce a stable snapshot within the timeout]\`) + expect(aborted).toMatchInlineSnapshot(\`true\`) + }) + " + `) +}) diff --git a/test/unit/test/expect-poll.test.ts b/test/unit/test/expect-poll.test.ts index 58eada44c..6b44ae664 100644 --- a/test/unit/test/expect-poll.test.ts +++ b/test/unit/test/expect-poll.test.ts @@ -1,4 +1,4 @@ -import { chai, expect, test, vi } from 'vitest' +import { expect, test, vi } from 'vitest' test('simple usage', async () => { await expect.poll(() => false).toBe(false) @@ -128,43 +128,49 @@ test('custom message', async () => { ).rejects.toMatchInlineSnapshot(`[AssertionError: custom: expected 1 to be 2 // Object.is equality]`) }) -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.toThrow() - fn.mock.results.forEach((result, index) => { - const isLastCall = index === fn.mock.results.length - 1 - expect(result.value).toBe(isLastCall ? true : undefined) - }) +test('unresolved function', async () => { + let aborted = false + await expect( + expect + .poll( + async ({ signal }) => { + signal.addEventListener('abort', () => { + aborted = true + }) + await new Promise(resolve => setTimeout(resolve, 500)) + return 'ok' + }, + { timeout: 50 }, + ) + .toBe('ok'), + ).rejects.toMatchInlineSnapshot(`[Error: expect.poll() function didn't resolve in time.]`) + expect(aborted).toBe(true) }) -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 +test('unresolved assertion', async () => { + expect.extend({ + toTestSlow: async () => { + await new Promise(resolve => setTimeout(resolve, 500)) + return { + pass: true, + message: () => 'ok', + } + }, }) - 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.toThrow(expect.objectContaining({ - // makes sure cause message reflects the last attempt value - message: 'expected 3 to be 1 // Object.is equality', - cause: expect.objectContaining({ - message: 'Matcher did not succeed in time.', - }), - })) + let aborted = false + await expect( + ( + expect.poll( + async ({ signal }) => { + signal.addEventListener('abort', () => { + aborted = true + }) + return 'ok' + }, + { timeout: 50 }, + ) as any + ).toTestSlow(), + ).rejects.toMatchInlineSnapshot(`[Error: expect.poll() assertion didn't resolve in time.]`) + expect(aborted).toBe(true) }) -- 2.51.2