From de19b8d7ce2c8a4b0e38ec1dbb24dd1afb9c0f3f Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 2 Nov 2026 16:45:28 -0800 Subject: [PATCH] test(integration): await completed callbacks and observe locks directly (cherry picked from commit 85e698a1ec6c75c32a56cab3c77d18ada772e5e4) --- .../test_team_delete_member_add_race.py | 38 +++++++++---------- .../test_logging_only_scope_config.py | 3 ++ 2 files changed, 22 insertions(+), 19 deletions(-) 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_logging_only_scope_config.py b/tests/integration/observability/test_logging_only_scope_config.py index fe255613316..212d8fd691c 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 ( _provider_response, _response_body_without_ids, _spend_row_for_call_id, + _spend_row_for_response_id, _spend_rows, wire_server, ) @@ -398,6 +399,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 = ( @@ -425,6 +427,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 = (