diff --git a/tests/e2e/quota_management/spend_tracking/test_cost_headers_e2e.py b/tests/e2e/quota_management/spend_tracking/test_cost_headers_e2e.py index a455f9f0db4..abc321ccde8 100644 --- a/tests/e2e/quota_management/spend_tracking/test_cost_headers_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_cost_headers_e2e.py @@ -102,11 +102,8 @@ class TestCostHeaders: return response return None - measured: StreamingResponse | None = None - for _ in range(CACHE_ATTEMPTS): - measured = prime_then_reread() - if measured is not None: - break + rounds = (prime_then_reread() for _ in range(CACHE_ATTEMPTS)) + measured = next((response for response in rounds if response is not None), None) if measured is None: pytest.fail( f"no cache read landed across {CACHE_ATTEMPTS} prime rounds of " diff --git a/tests/e2e/ui/tests/settings/routerSettings.spec.ts b/tests/e2e/ui/tests/settings/routerSettings.spec.ts index ada8e99e5c1..cd64e6e4453 100644 --- a/tests/e2e/ui/tests/settings/routerSettings.spec.ts +++ b/tests/e2e/ui/tests/settings/routerSettings.spec.ts @@ -139,24 +139,18 @@ async function patchRouterSettings( } /** - * Requires a consecutive streak because a single reply only proves the one replica that - * served it has reloaded, not the sibling still answering from the pre-update config. + * Spreads its samples across more than one reload cycle: a single reply only proves the one + * replica that served it has reloaded, not the sibling still on the pre-update config. */ -async function pollUntilSettled( - probe: () => Promise, - matches: (status: number) => boolean, - message: string, -): Promise { - let streak = 0; - await expect - .poll( - async () => { - streak = matches(await probe()) ? streak + 1 : 0; - return streak; - }, - { timeout: SETTLE_TIMEOUT_MS, intervals: [SETTLE_INTERVAL_MS], message }, - ) - .toBeGreaterThanOrEqual(SETTLE_PROBES); +async function sampleStatuses(probe: () => Promise): Promise { + return Array.from({ length: SETTLE_PROBES }).reduce>( + async (taken, _unused, index) => { + const sofar = await taken; + if (index > 0) await new Promise((resolve) => setTimeout(resolve, SETTLE_INTERVAL_MS)); + return [...sofar, await probe()]; + }, + Promise.resolve([]), + ); } test.describe("Router Settings - Loadbalancing", () => { @@ -289,19 +283,24 @@ test.describe("Router Settings - Fallbacks serve the request", () => { }) ).status(); - // The control: it proves the reply below could only have come from the fallback. - await pollUntilSettled( - chatStatus, - (status) => status >= 400, - "broken primary unexpectedly succeeded on its own", - ); + // The control: every replica must reject, or the reply below could have come from one + // that was still serving a fallback left behind by an earlier attempt. + await expect + .poll(async () => (await sampleStatuses(chatStatus)).every((status) => status >= 400), { + timeout: SETTLE_TIMEOUT_MS, + message: "broken primary unexpectedly succeeded on its own", + }) + .toBe(true); await patchRouterSettings(request, { fallbacks: [{ [BROKEN_PRIMARY]: [PRIMARY] }], } as Partial>); - // Same call now succeeds, served by the fallback model. - await pollUntilSettled(chatStatus, (status) => status === 200, "fallback never took effect"); + // One success is the whole claim here, so this waits for a first sighting rather than + // for every replica: demanding a streak would also assert a fallback hit rate. + await expect + .poll(chatStatus, { timeout: SETTLE_TIMEOUT_MS, message: "fallback never took effect" }) + .toBe(200); // And the playground renders a reply for a model whose own upstream is down. await openPlayground(page);