diff --git a/tests/e2e/migrations/database.py b/tests/e2e/migrations/database.py index 9bf03a216c4..34d94025a6e 100644 --- a/tests/e2e/migrations/database.py +++ b/tests/e2e/migrations/database.py @@ -89,6 +89,15 @@ class Database: (key >> 32, key & 0xFFFFFFFF), ) + @step("List the database sessions holding an advisory lock") + def holders(self, key: int = GATE_KEY) -> tuple[tuple[Scalar, ...], ...]: + return self.query( + "SELECT pid FROM pg_locks WHERE locktype = 'advisory' AND granted " + "AND database = (SELECT oid FROM pg_database WHERE datname = current_database()) " + "AND classid = %s AND objid = %s ORDER BY pid", + (key >> 32, key & 0xFFFFFFFF), + ) + @step("List replica identities waiting on any of the advisory locks") def waiting_identities(self, keys: tuple[int, ...]) -> frozenset[str]: if not keys: diff --git a/tests/e2e/migrations/test_startup.py b/tests/e2e/migrations/test_startup.py index 4a43ed4c0e3..5ec8d7690f2 100644 --- a/tests/e2e/migrations/test_startup.py +++ b/tests/e2e/migrations/test_startup.py @@ -14,7 +14,7 @@ from .checks import ( start_replicas, ) from .containers import Containers, failed, ready, until, waiting -from .database import GATE_KEY, PRISMA_LOCK, Database, Databases, restricted_user +from .database import COORDINATOR_LOCK, GATE_KEY, PRISMA_LOCK, Database, Databases, restricted_user from .startup_models import Migration from e2e_metadata import Domain, Subject, meta @@ -97,8 +97,11 @@ class TestMigrationStartup: with containers.start( database, (COMPLETE,), environment={"LITELLM_MIGRATION_LOCK_TIMEOUT": "12"} ) as replica: - until("Prisma lock contention", lambda: bool(database.blocked(PRISMA_LOCK))) - failed((replica,), "Timed out waiting for") + until( + "a replica holding the coordinator lock while Prisma's lock is held", + lambda: bool(database.holders(COORDINATOR_LOCK)), + ) + failed((replica,), "Timed out waiting for Prisma's lock") assert database.history() == history assert not database.exists("migration_effect") with containers.start(database, (COMPLETE,)) as restarted: diff --git a/tests/integration/cost_calculation/cases/misc.py b/tests/integration/cost_calculation/cases/misc.py index 389e552e44c..d17ab7a3d31 100644 --- a/tests/integration/cost_calculation/cases/misc.py +++ b/tests/integration/cost_calculation/cases/misc.py @@ -599,7 +599,13 @@ XAI_GROK_5_PROVIDER_REPORTED_COST: Final = CostTrackingTestCase( "usage": {"prompt_tokens": 1840, "completion_tokens": 412, "total_tokens": 2252, "cost": 0.0421}, }, ), - expected=ExactExpected(spend=0.0421, input_cost=0.0, output_cost=0.0421, prompt_tokens=1840, completion_tokens=412), + expected=ExactExpected( + spend=0.0421, + input_cost=0.029078078, + output_cost=0.013021922, + prompt_tokens=1840, + completion_tokens=412, + ), )