From 46913be92b785f7cb3cbc0cdc9dde354572cbb2f Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 15 Oct 2026 13:48:51 -0700 Subject: [PATCH] test(integration): wait on the right signal in five flaky integration contracts (#46737) * test(integration): wait on the right signal in five flaky integration contracts The Prometheus shared-label tests waited for the latency metric and then read the request counter, which the logger writes after an await. They now wait for the counter The auto-router session poll returned on the first 200 while the caller expected two turns, so it now waits for the expected turn count The S3 queue cap test let its failing sink recover on a 15s clock that could expire mid-flush. The sink now stays down until a flush fails after every event is queued K3 normally takes 168-175s against a 180s limit, so the limit is now 300s The MCP env var browser spec reused an idle keep-alive socket that the proxy closes after 5s, so its API calls now send Connection: close * test(integration): add until_asserted so a wait polls the exact asserts it guards eventually() takes a predicate separate from the asserts that follow it, which lets a test wait on one value and then assert on another written later. until_asserted reruns a block of asserts until they all pass and re-raises the last failure on timeout The Prometheus shared-label tests and the auto-router session check now use it --- .../mcpUserEnvVars.spec.ts | 2 +- tests/integration/_support/client.py | 11 ++ .../test_logging_only_scope_chaos.py | 2 +- .../test_prometheus_shared_custom_labels.py | 162 +++++++++--------- .../observability/test_s3_v2_upload_fanout.py | 9 +- .../routing/test_auto_router_management.py | 69 ++++---- 6 files changed, 138 insertions(+), 117 deletions(-) diff --git a/tests/e2e/ui/tests/integrationCritical/mcpUserEnvVars.spec.ts b/tests/e2e/ui/tests/integrationCritical/mcpUserEnvVars.spec.ts index 11cc1c948a3..fee09b409c7 100644 --- a/tests/e2e/ui/tests/integrationCritical/mcpUserEnvVars.spec.ts +++ b/tests/e2e/ui/tests/integrationCritical/mcpUserEnvVars.spec.ts @@ -12,7 +12,7 @@ import { captureRequestBody } from "../../helpers/roundTrip"; import { masterKey } from "../../helpers/traffic"; const master = masterKey(); -const headers = { Authorization: `Bearer ${master}` }; +const headers = { Authorization: `Bearer ${master}`, Connection: "close" }; const TOKEN = "USER_TOKEN"; type EnvVarStatus = { diff --git a/tests/integration/_support/client.py b/tests/integration/_support/client.py index b22c1b1da3a..053318c2451 100644 --- a/tests/integration/_support/client.py +++ b/tests/integration/_support/client.py @@ -57,6 +57,17 @@ def eventually( time.sleep(0.1) +def until_asserted(check: Callable[[], T], *, seconds: float = 10) -> T: + deadline: Final = time.monotonic() + seconds + while True: + try: + return check() + except AssertionError: + if time.monotonic() >= deadline: + raise + time.sleep(0.1) + + def held(read: Callable[[], T], satisfied: Callable[[T], bool], *, holding: float, seconds: float = 10) -> T: """`eventually`, and then `satisfied` must hold on every read for `holding` more seconds; a miss restarts the hold.""" deadline: Final = time.monotonic() + seconds diff --git a/tests/integration/observability/test_logging_only_scope_chaos.py b/tests/integration/observability/test_logging_only_scope_chaos.py index afcaab91b5f..a65e56b930b 100644 --- a/tests/integration/observability/test_logging_only_scope_chaos.py +++ b/tests/integration/observability/test_logging_only_scope_chaos.py @@ -333,7 +333,7 @@ def test_K2_policy_edge_delay_does_not_delay_concurrent_callers(gateway: Gateway ) == tuple((identity, "logging_only", "success") for _ in expected_directions), entries -@pytest.mark.timeout(180) +@pytest.mark.timeout(300) def test_K3_two_worker_sigkill_checks_post_kill_spend_rows( gateway: Gateway, tmp_path: Path, diff --git a/tests/integration/observability/test_prometheus_shared_custom_labels.py b/tests/integration/observability/test_prometheus_shared_custom_labels.py index 4c4b96c9509..c125c08e829 100644 --- a/tests/integration/observability/test_prometheus_shared_custom_labels.py +++ b/tests/integration/observability/test_prometheus_shared_custom_labels.py @@ -17,7 +17,7 @@ import pytest import yaml from integration._support import claude_code as cc from integration._support import responses_vendor as rv -from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment +from integration._support.client import Gateway, Scenario, gateway_from_environment, until_asserted from integration._support.process import owned_proxy from integration._support.prometheus_series import Sample, scrape from integration._support.wire import Reply, Request, Wire, wire_server @@ -233,26 +233,27 @@ def test_a_project_key_call_on_every_endpoint_lands_in_the_request_latency_and_p assert _send(rig.gateway, call, owner.key) == 200 forwarded: Final = rig.upstream.drain() assert [_forwarded(request) for request in forwarded] == [call.upstream], forwarded - samples: Final = eventually( - lambda: scrape(rig.gateway), - lambda found: ( - _total(found, _LATENCY_COUNT, hashed_api_key=owner.hashed) == 1 - and len(_project_gauges(found, owner.project)) == 4 - ), - seconds=30, - ) - requests: Final = _series(samples, _TOTAL_REQUESTS, hashed_api_key=owner.hashed) - assert [(sample.labels["route"], sample.labels["status_code"], sample.value) for sample in requests] == [ - (call.route, "200", 1.0) - ], requests - assert [_shared_labels(sample) for sample in requests] == [(owner.alias, owner.team, call.model, "None")] - gauges: Final = _project_gauges(samples, owner.project) - assert {(sample.name, sample.labels["rate_limit_type"], int(sample.value)) for sample in gauges} >= { - (_ALLOWED, "requests", _RPM), - (_USED, "requests", 1), - (_ALLOWED, "tokens", _TPM), - }, gauges - assert {_shared_labels(sample) for sample in gauges} == {(owner.alias, owner.team, call.model, "None")}, gauges + + def landed() -> None: + samples: Final = scrape(rig.gateway) + assert _total(samples, _LATENCY_COUNT, hashed_api_key=owner.hashed) == 1 + requests: Final = _series(samples, _TOTAL_REQUESTS, hashed_api_key=owner.hashed) + assert [(sample.labels["route"], sample.labels["status_code"], sample.value) for sample in requests] == [ + (call.route, "200", 1.0) + ], requests + assert [_shared_labels(sample) for sample in requests] == [(owner.alias, owner.team, call.model, "None")] + gauges: Final = _project_gauges(samples, owner.project) + assert len(gauges) == 4, gauges + assert {(sample.name, sample.labels["rate_limit_type"], int(sample.value)) for sample in gauges} >= { + (_ALLOWED, "requests", _RPM), + (_USED, "requests", 1), + (_ALLOWED, "tokens", _TPM), + }, gauges + assert {_shared_labels(sample) for sample in gauges} == {(owner.alias, owner.team, call.model, "None")}, ( + gauges + ) + + until_asserted(landed, seconds=30) def test_a_failed_project_key_call_counts_as_one_failed_request_with_each_shared_label_once(rig: _Rig) -> None: @@ -264,15 +265,19 @@ def test_a_failed_project_key_call_counts_as_one_failed_request_with_each_shared assert status >= 500, status forwarded: Final = rig.upstream.drain() assert len(forwarded) == 1, forwarded - samples: Final = eventually( - lambda: scrape(rig.gateway), - lambda found: _total(found, _FAILED_REQUESTS, hashed_api_key=owner.hashed) == 1, - seconds=30, - ) - failed: Final = _series(samples, _FAILED_REQUESTS, hashed_api_key=owner.hashed) - assert [_shared_labels(sample) for sample in failed] == [("None", owner.team, _DOWN, "None")], failed - requests: Final = _series(samples, _TOTAL_REQUESTS, hashed_api_key=owner.hashed) - assert [(sample.labels["status_code"], sample.value) for sample in requests] == [(str(status), 1.0)], requests + + def landed() -> None: + samples: Final = scrape(rig.gateway) + failed: Final = _series(samples, _FAILED_REQUESTS, hashed_api_key=owner.hashed) + assert [(_shared_labels(sample), sample.value) for sample in failed] == [ + (("None", owner.team, _DOWN, "None"), 1.0) + ], failed + requests: Final = _series(samples, _TOTAL_REQUESTS, hashed_api_key=owner.hashed) + assert [(sample.labels["status_code"], sample.value) for sample in requests] == [(str(status), 1.0)], ( + requests + ) + + until_asserted(landed, seconds=30) def test_request_metadata_fills_the_shared_labels_on_the_request_and_project_metrics(rig: _Rig) -> None: @@ -288,19 +293,17 @@ def test_request_metadata_fills_the_shared_labels_on_the_request_and_project_met assert _send(rig.gateway, call, owner.key, metadata=metadata) == 200 forwarded: Final = rig.upstream.drain() assert len(forwarded) == 1, forwarded - samples: Final = eventually( - lambda: scrape(rig.gateway), - lambda found: ( - _total(found, _TOTAL_REQUESTS, hashed_api_key=owner.hashed) == 1 - and len(_project_gauges(found, owner.project)) == 4 - ), - seconds=30, - ) expected: Final = ("caller-alias", "caller-team", _GPT, "eng") - requests: Final = _series(samples, _TOTAL_REQUESTS, hashed_api_key=owner.hashed) - assert [_shared_labels(sample) for sample in requests] == [expected], requests - gauges: Final = _project_gauges(samples, owner.project) - assert {_shared_labels(sample) for sample in gauges} == {expected}, gauges + + def landed() -> None: + samples: Final = scrape(rig.gateway) + requests: Final = _series(samples, _TOTAL_REQUESTS, hashed_api_key=owner.hashed) + assert [(_shared_labels(sample), sample.value) for sample in requests] == [(expected, 1.0)], requests + gauges: Final = _project_gauges(samples, owner.project) + assert len(gauges) == 4, gauges + assert {_shared_labels(sample) for sample in gauges} == {expected}, gauges + + until_asserted(landed, seconds=30) def test_a_key_outside_any_team_or_project_still_counts_its_request(rig: _Rig) -> None: @@ -311,15 +314,16 @@ def test_a_key_outside_any_team_or_project_still_counts_its_request(rig: _Rig) - assert _send(rig.gateway, _Call("/v1/chat/completions", _GPT, stream=False), key) == 200 forwarded: Final = rig.upstream.drain() assert len(forwarded) == 1, forwarded - samples: Final = eventually( - lambda: scrape(rig.gateway), - lambda found: _total(found, _LATENCY_COUNT, hashed_api_key=hashed) == 1, - seconds=30, - ) - requests: Final = _series(samples, _TOTAL_REQUESTS, hashed_api_key=hashed) - assert [(_shared_labels(sample), sample.value) for sample in requests] == [ - (("None", "None", _GPT, "None"), 1.0) - ], requests + + def landed() -> None: + samples: Final = scrape(rig.gateway) + assert _total(samples, _LATENCY_COUNT, hashed_api_key=hashed) == 1 + requests: Final = _series(samples, _TOTAL_REQUESTS, hashed_api_key=hashed) + assert [(_shared_labels(sample), sample.value) for sample in requests] == [ + (("None", "None", _GPT, "None"), 1.0) + ], requests + + until_asserted(landed, seconds=30) def test_a_mixed_burst_through_an_upstream_outage_counts_every_request_once_and_never_fails_the_callback( @@ -346,24 +350,21 @@ def test_a_mixed_burst_through_an_upstream_outage_counts_every_request_once_and_ assert all(status >= 500 for _, status in down), down rig.upstream.drain() observed: Final = Counter((route, str(status)) for route, status in healthy + down + recovered) - samples: Final = eventually( - lambda: scrape(rig.gateway), - lambda found: ( - sum(sample.value for sample in _series(found, _TOTAL_REQUESTS, hashed_api_key=owner.hashed)) - == len(calls) * 3 - ), - seconds=60, - ) - counted: Final = Counter( - { - (sample.labels["route"], sample.labels["status_code"]): int(sample.value) - for sample in _series(samples, _TOTAL_REQUESTS, hashed_api_key=owner.hashed) - } - ) - assert counted == observed, (counted, observed) - assert _total(samples, _LATENCY_COUNT, hashed_api_key=owner.hashed) == len(healthy) + len(recovered) - assert _total(samples, _FAILED_REQUESTS, hashed_api_key=owner.hashed) == len(down) - assert _total(samples, _CALLBACK_FAILURES, callback_name="PrometheusLogger") == before + + def landed() -> None: + samples: Final = scrape(rig.gateway) + counted: Final = Counter( + { + (sample.labels["route"], sample.labels["status_code"]): int(sample.value) + for sample in _series(samples, _TOTAL_REQUESTS, hashed_api_key=owner.hashed) + } + ) + assert counted == observed, (counted, observed) + assert _total(samples, _LATENCY_COUNT, hashed_api_key=owner.hashed) == len(healthy) + len(recovered) + assert _total(samples, _FAILED_REQUESTS, hashed_api_key=owner.hashed) == len(down) + assert _total(samples, _CALLBACK_FAILURES, callback_name="PrometheusLogger") == before + + until_asserted(landed, seconds=60) @pytest.mark.parametrize( @@ -381,15 +382,14 @@ def test_a_project_key_call_under_other_label_configurations_keeps_one_project_a assert _send(owned.gateway, _Call("/v1/chat/completions", _GPT, stream=False), owner.key) == 200 forwarded: Final = owned.upstream.drain() assert len(forwarded) == 1, forwarded - samples: Final = eventually( - lambda: scrape(owned.gateway), - lambda found: ( - _total(found, _LATENCY_COUNT, hashed_api_key=owner.hashed) == 1 - and len(_project_gauges(found, owner.project)) == 4 - ), - seconds=30, - ) - requests: Final = _series(samples, _TOTAL_REQUESTS, hashed_api_key=owner.hashed) - assert [(sample.labels["project_alias"], sample.value) for sample in requests] == [(owner.alias, 1.0)] - gauges: Final = _project_gauges(samples, owner.project) - assert {sample.labels["project_alias"] for sample in gauges} == {owner.alias}, gauges + + def landed() -> None: + samples: Final = scrape(owned.gateway) + assert _total(samples, _LATENCY_COUNT, hashed_api_key=owner.hashed) == 1 + requests: Final = _series(samples, _TOTAL_REQUESTS, hashed_api_key=owner.hashed) + assert [(sample.labels["project_alias"], sample.value) for sample in requests] == [(owner.alias, 1.0)] + gauges: Final = _project_gauges(samples, owner.project) + assert len(gauges) == 4, gauges + assert {sample.labels["project_alias"] for sample in gauges} == {owner.alias}, gauges + + until_asserted(landed, seconds=30) diff --git a/tests/integration/observability/test_s3_v2_upload_fanout.py b/tests/integration/observability/test_s3_v2_upload_fanout.py index 54e35d993e6..01c1f563be4 100644 --- a/tests/integration/observability/test_s3_v2_upload_fanout.py +++ b/tests/integration/observability/test_s3_v2_upload_fanout.py @@ -1,4 +1,5 @@ import json +import math import re import threading import time @@ -518,16 +519,20 @@ def test_s3_v2_failing_sink_trims_the_oldest_failed_events_past_the_queue_cap(ga ): model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key") key: Final = scenario.key(models=[model]) - sink.fail_until = time.time() + 15 + sink.fail_until = math.inf _send(owned.gateway, model, key, f"{marker}-probe") eventually(lambda: owned.log.read_text(), lambda text: "S3BatchUploadError" in text, seconds=30) for index in range(24): _send(owned.gateway, model, key, f"{marker}-{index}") + sent: Final = len(owned.log.read_text()) eventually( lambda: owned.log.read_text(), - lambda text: "after a failed flush, dropped" in text, + lambda text: ( + "after a failed flush, dropped" in text and text[sent:].count("async_send_batch raised") >= 2 + ), seconds=30, ) + sink.fail_until = time.time() payloads: Final = collect_payloads(sink, 4, seconds=90) landed: Final = frozenset(payload["id"] for payload in payloads) assert sum(1 for r in provider.drain() if r.method == "POST") == 25 diff --git a/tests/integration/routing/test_auto_router_management.py b/tests/integration/routing/test_auto_router_management.py index 7d46ac4e702..5aed328ba7b 100644 --- a/tests/integration/routing/test_auto_router_management.py +++ b/tests/integration/routing/test_auto_router_management.py @@ -8,7 +8,7 @@ from typing import Final import httpx import pytest import yaml -from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value +from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value, until_asserted from integration._support.database import read_rows, scratch_database from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, Wire, wire_server @@ -393,37 +393,42 @@ def test_auto_router_session_tracks_key_scoped_spend_and_bounds_long_session_ids {"type": "text", "text": "synthetic tier response"} ], messages_second_response.text - messages_session: Final = _session_response(candidate, key, messages_session_id) - assert JSON_OBJECT.validate_json(messages_session.content) == { - "session_id": messages_session_id, - "router_name": router_name, - "router_type": "complexity", - "turns": 2, - "last_model": "openai/integration-medium", - "spend": pytest.approx(0.000045), - "savings_estimated_turns": 2, - "savings_estimated_actual_spend": pytest.approx(0.000045), - "saved_spend": pytest.approx(0.000255), - "baseline_spend": pytest.approx(0.0003), - "savings_estimated_baseline_spend": pytest.approx(0.0003), - "baseline_model": "openai/integration-baseline", - "baseline_models": {"openai/integration-baseline": 2}, - }, messages_session.text - assert _rollup_row(key, database_url, router_name, messages_session_id) == { - "session_id": messages_session_id, - "router_name": router_name, - "router_type": "complexity", - "last_model": "openai/integration-medium", - "turns": 2, - "spend": pytest.approx(0.000045), - "saved_spend": pytest.approx(0.000255), - "savings_estimated_turns": 2, - "savings_estimated_actual_spend": pytest.approx(0.000045), - "savings_estimated_saved_spend": pytest.approx(0.000255), - "savings_estimated_baseline_models": {"openai/integration-baseline": 2}, - "baseline_models": {"openai/integration-baseline": 2}, - "tier_turns": {"SIMPLE": 1, "MEDIUM": 1}, - } + def messages_session_settled() -> None: + session: Final = candidate.request( + "GET", "/auto_router/session", params={"session_id": messages_session_id}, key=key + ) + assert JSON_OBJECT.validate_json(session.content) == { + "session_id": messages_session_id, + "router_name": router_name, + "router_type": "complexity", + "turns": 2, + "last_model": "openai/integration-medium", + "spend": pytest.approx(0.000045), + "savings_estimated_turns": 2, + "savings_estimated_actual_spend": pytest.approx(0.000045), + "saved_spend": pytest.approx(0.000255), + "baseline_spend": pytest.approx(0.0003), + "savings_estimated_baseline_spend": pytest.approx(0.0003), + "baseline_model": "openai/integration-baseline", + "baseline_models": {"openai/integration-baseline": 2}, + }, session.text + assert _rollup_row(key, database_url, router_name, messages_session_id) == { + "session_id": messages_session_id, + "router_name": router_name, + "router_type": "complexity", + "last_model": "openai/integration-medium", + "turns": 2, + "spend": pytest.approx(0.000045), + "saved_spend": pytest.approx(0.000255), + "savings_estimated_turns": 2, + "savings_estimated_actual_spend": pytest.approx(0.000045), + "savings_estimated_saved_spend": pytest.approx(0.000255), + "savings_estimated_baseline_models": {"openai/integration-baseline": 2}, + "baseline_models": {"openai/integration-baseline": 2}, + "tier_turns": {"SIMPLE": 1, "MEDIUM": 1}, + } + + until_asserted(messages_session_settled, seconds=70) responses_response: Final = candidate.client.post( "/v1/responses",