diff --git a/tests/integration/management/test_team_delete_member_add_race.py b/tests/integration/management/test_team_delete_member_add_race.py index e1f2a933ec3..c558264a650 100644 --- a/tests/integration/management/test_team_delete_member_add_race.py +++ b/tests/integration/management/test_team_delete_member_add_race.py @@ -22,8 +22,12 @@ import time import uuid from contextlib import asynccontextmanager from datetime import timedelta +from typing import Final from unittest.mock import MagicMock +from psycopg import AsyncConnection +from psycopg.rows import dict_row + import pytest from fastapi import HTTPException @@ -48,7 +52,7 @@ _HELD_LOCK_KEY_SQL = ( _LOCK_WAITER_SQL = ( "SELECT count(*)::int AS waiters FROM pg_locks " "WHERE locktype = 'advisory' AND NOT granted " - "AND classid::bigint = $1 AND objid::bigint = $2" + "AND classid::bigint = %s AND objid::bigint = %s" ) _LOCK_WAIT_TIMEOUT_SECONDS = 20.0 _LOCK_POLL_SECONDS = 0.01 @@ -66,21 +70,17 @@ async def _hold_team_lock(held, team_id: str) -> tuple[int, int]: return rows[0]["classid"], rows[0]["objid"] -async def _await_lock_contention(watcher, lock_key: tuple[int, int], task, what: str) -> None: - """Block until Postgres reports `task` queued behind the held lock. - - This is the assertion that the endpoint serializes on the team's advisory lock, and it - is what a fixed sleep was standing in for: the endpoint is only provably waiting once a - non-granted advisory lock on the same key exists. `watcher` must be a connection that is - not itself blocked, so it can observe the queue.""" - classid, objid = lock_key - deadline = time.monotonic() + _LOCK_WAIT_TIMEOUT_SECONDS - while time.monotonic() < deadline: - if task.done(): - raise AssertionError(f"{what} returned without waiting on the team's advisory lock") from task.exception() - if (await watcher.query_raw(_LOCK_WAITER_SQL, classid, objid))[0]["waiters"]: - return - await asyncio.sleep(_LOCK_POLL_SECONDS) +async def _await_lock_contention(lock_key: tuple[int, int], task: asyncio.Task[object], what: str) -> None: + deadline: Final = time.monotonic() + _LOCK_WAIT_TIMEOUT_SECONDS + async with await AsyncConnection.connect(os.environ["DATABASE_URL"], autocommit=True, row_factory=dict_row) as watcher: + while time.monotonic() < deadline: + if task.done(): + raise AssertionError(f"{what} returned without waiting on the team's advisory lock") from task.exception() + cursor: Final = await watcher.execute(_LOCK_WAITER_SQL, lock_key) + observed: Final = await cursor.fetchone() + if observed is not None and observed["waiters"]: + return + await asyncio.sleep(_LOCK_POLL_SECONDS) raise AssertionError(f"{what} never queued on the team's advisory lock within {_LOCK_WAIT_TIMEOUT_SECONDS}s") @@ -171,7 +171,7 @@ async def test_member_add_blocked_by_delete_writes_no_dangling_reference(): async with blocker.tx(timeout=timedelta(seconds=30)) as held: lock_key = await _hold_team_lock(held, team_id) task = asyncio.create_task(add_member()) - await _await_lock_contention(db, lock_key, task, "member_add") + await _await_lock_contention(lock_key, task, "member_add") await held.execute_raw(_DELETE_TEAM, team_id) @@ -234,7 +234,7 @@ async def test_member_delete_blocked_by_member_add_removes_from_the_fresh_roster async with blocker.tx(timeout=timedelta(seconds=30)) as held: lock_key = await _hold_team_lock(held, team_id) task = asyncio.create_task(run_delete()) - await _await_lock_contention(db, lock_key, task, "member_delete") + await _await_lock_contention(lock_key, task, "member_delete") await held.litellm_teamtable.update( where={"team_id": team_id}, @@ -311,7 +311,7 @@ async def test_delete_blocked_by_member_add_sweeps_the_fresh_reference(): async with blocker.tx(timeout=timedelta(seconds=30)) as held: lock_key = await _hold_team_lock(held, team_id) task = asyncio.create_task(run_delete()) - await _await_lock_contention(db, lock_key, task, "delete_team") + await _await_lock_contention(lock_key, task, "delete_team") await held.litellm_usertable.upsert( where={"user_id": user_id}, diff --git a/tests/integration/observability/test_cache_hit_guardrail_metrics.py b/tests/integration/observability/test_cache_hit_guardrail_metrics.py index 13cd2133d4f..3abb2a7e2be 100644 --- a/tests/integration/observability/test_cache_hit_guardrail_metrics.py +++ b/tests/integration/observability/test_cache_hit_guardrail_metrics.py @@ -424,7 +424,7 @@ def test_cache_hit_post_call_reject_keeps_deployment_labels_anthropic(gateway: G raised = True assert raised, "cache-hit post_call guardrail did not reject /v1/messages" assert rig.provider.received.qsize() == 1, rig.provider.drain() - _expect_metrics(rig, 1, 0, api_provider="anthropic", pf_id="None") + _expect_metrics(rig, 1, 0, api_provider="anthropic") def test_cache_hit_post_call_reject_keeps_deployment_labels_responses(gateway: Gateway, tmp_path: Path) -> None: @@ -441,7 +441,7 @@ def test_cache_hit_post_call_reject_keeps_deployment_labels_responses(gateway: G ) assert reject.status_code == 400, reject.text assert rig.provider.received.qsize() == 1, rig.provider.drain() - _expect_metrics(rig, 1, 0, pf_id="None") + _expect_metrics(rig, 1, 0) _assert_spend(_call_id(reject), rig) diff --git a/tests/integration/observability/test_logging_only_scope_config.py b/tests/integration/observability/test_logging_only_scope_config.py index ea524c15c92..5a37e779ef3 100644 --- a/tests/integration/observability/test_logging_only_scope_config.py +++ b/tests/integration/observability/test_logging_only_scope_config.py @@ -35,6 +35,7 @@ from _logging_only_scope_support import ( _reset_upstream, _response_body_without_ids, _spend_row_for_call_id, + _spend_row_for_response_id, _spend_rows, _upstream_requests, wire_server, @@ -404,6 +405,7 @@ def test_presidio_scope_matches_no_scope_behavior(gateway: Gateway, tmp_path: Pa bool, seconds=30, ) + _spend_row_for_response_id(no_scope_result.response_id) no_scope_analyzer_calls: Final = no_scope_analyzer.drain() no_scope_anonymizer_calls: Final = no_scope_anonymizer.drain() no_scope_call_counts: Final = ( @@ -431,6 +433,7 @@ def test_presidio_scope_matches_no_scope_behavior(gateway: Gateway, tmp_path: Pa bool, seconds=30, ) + _spend_row_for_response_id(scoped_result.response_id) scoped_analyzer_calls: Final = scoped_analyzer.drain() scoped_anonymizer_calls: Final = scoped_anonymizer.drain() scoped_call_counts: Final = ( diff --git a/tests/integration/spend/test_normalized_error_long_message.py b/tests/integration/spend/test_normalized_error_long_message.py index a364327290f..23200c0450e 100644 --- a/tests/integration/spend/test_normalized_error_long_message.py +++ b/tests/integration/spend/test_normalized_error_long_message.py @@ -30,6 +30,7 @@ BENIGN_288KB_MODEL: Final = "x" * len(CRAFTED_MODEL) CRAFTED_SLOWDOWN_FACTOR: Final = 10 FAST_SECONDS: Final = 10.0 LIVELINESS_STALL_BUDGET_SECONDS: Final = 5.0 +BURST_STALL_BUDGET_SECONDS: Final = 5.0 ROW_SECONDS: Final = 70 CHAT: Final = "/v1/chat/completions" MESSAGES: Final = "/v1/messages" @@ -349,10 +350,16 @@ def _burst_call(client: httpx.Client, label: str, path: str, body: Mapping[str, def _burst( - client: httpx.Client, happy_model: str, happy_key: str, open_key: str, during: Callable[[], None] + client: httpx.Client, + happy_model: str, + happy_key: str, + open_key: str, + during: Callable[[], None], + *, + rejected_model: str = CRAFTED_MODEL, ) -> tuple[_BurstResult, ...]: crafted: Final = tuple( - (f"crafted-{path}-{index}", path, _body(path, CRAFTED_MODEL, uuid.uuid4().hex, stream=index % 2 == 1), open_key) + (f"crafted-{path}-{index}", path, _body(path, rejected_model, uuid.uuid4().hex, stream=index % 2 == 1), open_key) for path in (CHAT, MESSAGES, RESPONSES) for index in range(4) ) @@ -377,6 +384,18 @@ def _burst( return tuple(future.result() for future in first + second) +def _assert_burst_rejections_fast(results: tuple[_BurstResult, ...], control: tuple[_BurstResult, ...]) -> None: + crafted: Final = tuple(result for result in results if result.label.startswith("crafted") and result.status) + benign: Final = tuple(result for result in control if result.label.startswith("crafted")) + assert crafted and all(result.status == 400 for result in crafted), crafted + assert benign and all(result.status == 400 for result in benign), benign + assert_not_delayed_by( + (result.seconds for result in crafted), + control=(result.seconds for result in benign), + injected_delay=BURST_STALL_BUDGET_SECONDS, + ) + + def _assert_rows_land_exactly_once(results: tuple[_BurstResult, ...], prefix: str) -> None: landed: Final = tuple(result for result in results if result.label.startswith(prefix) and result.status == 200) assert landed, results @@ -405,20 +424,18 @@ def test_killing_one_worker_mid_burst_leaves_the_other_serving_crafted_and_happy ): model: Final = scenario.model(api_base=wire.url + "/v1", num_retries=0) key: Final = scenario.key(models=[model]) + open_key: Final = scenario.key() + control: Final = _burst(client, model, key, open_key, lambda: None, rejected_model=BENIGN_288KB_MODEL) + assert all(result.status == (400 if result.label.startswith("crafted") else 200) for result in control), control workers: Final = owned.worker_pids() assert len(workers) == 2, workers victim: Final = min(workers) with owned.replacing_worker(victim) as kill: - results: Final = _burst(client, model, key, scenario.key(), kill) + results: Final = _burst(client, model, key, open_key, kill) dropped: Final = tuple(result for result in results if result.status is None) assert len(dropped) < len(results), results - crafted: Final = tuple( - result for result in results if result.label.startswith("crafted") and result.status - ) - assert crafted and all( - result.status == 400 and result.seconds < FAST_SECONDS for result in crafted - ), crafted + _assert_burst_rejections_fast(results, control) late: Final = tuple(result for result in results if result.label.startswith("late")) assert all(result.status == 200 for result in late), late _assert_rows_land_exactly_once(results, "late") @@ -446,7 +463,11 @@ def test_upstream_returning_503_mid_burst_logs_every_failure_with_its_own_cluste ): model: Final = scenario.model(api_base=wire.url + "/v1", num_retries=0) key: Final = scenario.key(models=[model]) - results: Final = _burst(client, model, key, scenario.key(), lambda: None) + open_key: Final = scenario.key() + control: Final = _burst(client, model, key, open_key, lambda: None, rejected_model=BENIGN_288KB_MODEL) + assert all(result.status == (400 if result.label.startswith("crafted") else 503) for result in control), control + wire.drain() + results: Final = _burst(client, model, key, open_key, lambda: None) assert all(result.status is not None for result in results), results happy: Final = tuple(result for result in results if result.label.startswith("happy")) assert all(result.status == 503 for result in happy), happy @@ -456,8 +477,7 @@ def test_upstream_returning_503_mid_burst_logs_every_failure_with_its_own_cluste assert all(any(f"normalized error audit {result.label}" in body for body in seen) for result in happy + late), ( seen ) - crafted: Final = tuple(result for result in results if result.label.startswith("crafted")) - assert all(result.status == 400 and result.seconds < FAST_SECONDS for result in crafted), crafted + _assert_burst_rejections_fast(results, control) codes: Final = { result.label: _error_information(string_value(result.call_id))["normalized_error"] for result in results }