mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-05 07:37:50 +00:00
229f5a9745
Scope the Prisma migration lock test's advisory lock check to its own database, so a lock another xdist worker holds in its database no longer trips it
Keep in-process MCP peers' idle connections open for 30s instead of 1s, so the gateway's pooled HTTP client doesn't reuse a connection the peer is closing and hang the tool call
Delete database guardrails in the logging-only scope tests through DELETE /guardrails/{id} on the shared gateway instead of raw SQL, so a default-on guardrail stops applying to the next test on the same worker right away
Bound the wedged coordination Redis /customer/update check by the shortest Redis wait the proxy uses, like the burst check, so a slow runner can't fail it while a handler that waits on Redis still does
82 lines
3.6 KiB
Python
82 lines
3.6 KiB
Python
import os
|
|
import secrets
|
|
from typing import Final
|
|
|
|
import psycopg
|
|
import pytest
|
|
from litellm_proxy_extras.migration_lock import MIGRATION_LOCK_KEY, PRISMA_ADVISORY_LOCK_KEY, prisma_migration_lock
|
|
|
|
pytestmark = pytest.mark.skipif("DATABASE_URL" not in os.environ, reason="needs a live Postgres")
|
|
|
|
|
|
def _can_take_prismas_session_lock(connection: psycopg.Connection[tuple[object, ...]]) -> bool:
|
|
row: Final = connection.execute("SELECT pg_try_advisory_lock(%s)", (PRISMA_ADVISORY_LOCK_KEY,)).fetchone()
|
|
taken: Final = row is not None and row[0] is True
|
|
if taken:
|
|
connection.execute("SELECT pg_advisory_unlock(%s)", (PRISMA_ADVISORY_LOCK_KEY,))
|
|
return taken
|
|
|
|
|
|
def _granted_advisory_keys(connection: psycopg.Connection[tuple[object, ...]]) -> frozenset[int]:
|
|
rows: Final = connection.execute(
|
|
"SELECT (classid::bigint << 32) | objid::bigint FROM pg_locks WHERE locktype = 'advisory' AND granted "
|
|
"AND database = (SELECT oid FROM pg_database WHERE datname = current_database())"
|
|
).fetchall()
|
|
return frozenset(int(row[0]) for row in rows)
|
|
|
|
|
|
def _index_builds_concurrently_now(connection: psycopg.Connection[tuple[object, ...]]) -> bool:
|
|
table: Final = f'"lit_cic_{secrets.token_hex(4)}"'
|
|
connection.execute(f"CREATE TABLE {table} (a int)")
|
|
try:
|
|
connection.execute("SET lock_timeout = '5s'")
|
|
connection.execute(f"CREATE INDEX CONCURRENTLY ON {table} (a)")
|
|
except psycopg.errors.LockNotAvailable:
|
|
return False
|
|
finally:
|
|
connection.execute("RESET lock_timeout")
|
|
connection.execute(f"DROP TABLE {table}")
|
|
return True
|
|
|
|
|
|
def test_prisma_migration_lock_does_not_stall_a_concurrent_index_build():
|
|
with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as peer:
|
|
with prisma_migration_lock(os.environ["DATABASE_URL"]):
|
|
assert _index_builds_concurrently_now(peer)
|
|
|
|
|
|
def test_prisma_migration_lock_holds_prismas_key_on_a_pinned_connection_and_leaves_nothing_behind():
|
|
with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as peer:
|
|
assert _can_take_prismas_session_lock(peer)
|
|
|
|
with prisma_migration_lock(os.environ["DATABASE_URL"]):
|
|
held: Final = _granted_advisory_keys(peer)
|
|
assert PRISMA_ADVISORY_LOCK_KEY in held
|
|
assert MIGRATION_LOCK_KEY not in held
|
|
assert not _can_take_prismas_session_lock(peer)
|
|
|
|
assert _granted_advisory_keys(peer).isdisjoint({PRISMA_ADVISORY_LOCK_KEY, MIGRATION_LOCK_KEY})
|
|
assert _can_take_prismas_session_lock(peer)
|
|
|
|
|
|
def test_prisma_migration_lock_does_not_wait_for_the_coordinator_lock(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setenv("LITELLM_MIGRATION_LOCK_TIMEOUT", "3")
|
|
with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as peer:
|
|
peer.execute("SELECT pg_advisory_lock(%s)", (MIGRATION_LOCK_KEY,))
|
|
|
|
with prisma_migration_lock(os.environ["DATABASE_URL"]):
|
|
assert PRISMA_ADVISORY_LOCK_KEY in _granted_advisory_keys(peer)
|
|
|
|
|
|
def test_prisma_migration_lock_times_out_when_another_prisma_migration_holds_its_key(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setenv("LITELLM_MIGRATION_LOCK_TIMEOUT", "3")
|
|
with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as peer:
|
|
peer.execute("SELECT pg_advisory_lock(%s)", (PRISMA_ADVISORY_LOCK_KEY,))
|
|
|
|
with pytest.raises(RuntimeError, match="Timed out waiting for another Prisma migration"):
|
|
with prisma_migration_lock(os.environ["DATABASE_URL"]):
|
|
pytest.fail("Prisma migration lock was acquired while another migration held it")
|