From efa7557857f253dfe64a4bf870c81e09c096fec9 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 22 Feb 2026 03:42:14 +0000 Subject: [PATCH] test(policy-versioning-e2e): mock prisma tx in lifecycle Co-authored-by: Krish Dholakia --- .../policy_engine/test_policy_versioning_e2e.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/proxy/policy_engine/test_policy_versioning_e2e.py b/tests/test_litellm/proxy/policy_engine/test_policy_versioning_e2e.py index 5d6f3a05ae6..c79ebecf32a 100644 --- a/tests/test_litellm/proxy/policy_engine/test_policy_versioning_e2e.py +++ b/tests/test_litellm/proxy/policy_engine/test_policy_versioning_e2e.py @@ -99,9 +99,15 @@ async def test_full_lifecycle_create_draft_edit_publish_promote(): guardrails_add=["g1", "g2"], description="Draft v2", ) - prisma.db.litellm_policytable.find_first = AsyncMock(return_value=created_v1) - prisma.db.litellm_policytable.update_many = AsyncMock() - prisma.db.litellm_policytable.create = AsyncMock(return_value=v2_row) + # create_new_version uses an interactive transaction + tx = MagicMock() + prisma.db.tx = MagicMock() + prisma.db.tx.return_value.__aenter__ = AsyncMock(return_value=tx) + prisma.db.tx.return_value.__aexit__ = AsyncMock(return_value=None) + # find_first for (1) production lookup and (2) latest version lookup + tx.litellm_policytable.find_first = AsyncMock(side_effect=[created_v1, created_v1]) + tx.litellm_policytable.update_many = AsyncMock() + tx.litellm_policytable.create = AsyncMock(return_value=v2_row) draft_v2 = await registry.create_new_version( policy_name="lifecycle-policy", @@ -164,6 +170,7 @@ async def test_full_lifecycle_create_draft_edit_publish_promote(): # 5) Promote v2 to production (demote v1 to published, update registry) prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=v2_published) + prisma.db.litellm_policytable.find_first = AsyncMock(return_value=created_v1) prisma.db.litellm_policytable.update_many = AsyncMock() v2_production = _make_row( policy_id="v2-id",