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",