scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

test(proxy): inject a fake prisma floor instead of patching SpendCounterReseed.from_db

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-09-14 23:32:25 +00:00
parent e91c6ca789
commit 54f11b29a4
1 changed files with 22 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

38
tests/test_litellm/proxy/test_budget_reservation.py
View file

@ -2230,6 +2230,17 @@ class _ExpiringRedisCache:
return None
class _TeamMembershipFloorDb:
"""Stands in for `prisma_client.db`: only the team-membership row exists and its spend is the DB floor."""
def __init__(self, spend: float) -> None:
self.spend = spend
def __getattr__(self, table_name: str) -> SimpleNamespace:
row = SimpleNamespace(spend=self.spend) if table_name == "litellm_teammembership" else None
return SimpleNamespace(find_unique=AsyncMock(return_value=row))
@pytest.mark.asyncio
async def test_reconcile_after_redis_counter_expiry_keeps_request_cost_enforced(
spend_counter_state,
@ -2292,6 +2303,8 @@ async def test_reconcile_before_db_update_does_not_double_count_when_flush_lands
redis_cache = _ExpiringRedisCache()
counter_cache.redis_cache = redis_cache
counter_cache.in_memory_cache.set_cache(key=counter_key, value=0.6)
db_floor = _TeamMembershipFloorDb(spend=0.3)
ps.prisma_client = SimpleNamespace(db=db_floor)
reservation = {
"reserved_cost": 0.6,
@ -2307,27 +2320,20 @@ async def test_reconcile_before_db_update_does_not_double_count_when_flush_lands
"finalized": False,
}
with patch.object( # test-quality-ok: the reseed reads the DB floor through a Prisma client the test has no seam for
ps.SpendCounterReseed, "from_db", AsyncMock(return_value=0.3)
):
await reconcile_budget_reservation(
budget_reservation=reservation, actual_cost=0.05, finalize=False
)
await reconcile_budget_reservation(budget_reservation=reservation, actual_cost=0.05, finalize=False)
assert redis_cache.store[counter_key] == pytest.approx(0.35)
assert reservation["entries"][0]["applied_adjustment"] == pytest.approx(-0.55)
assert reservation["finalized"] is False
with patch.object( # test-quality-ok: the flush landing between the passes makes the DB floor include this request
ps.SpendCounterReseed, "from_db", AsyncMock(return_value=0.35)
):
await ps.increment_spend_counters(
token="key-flush",
team_id="team-flush",
user_id="user-flush",
response_cost=0.05,
budget_reservation=reservation,
)
db_floor.spend = 0.35
await ps.increment_spend_counters(
token="key-flush",
team_id="team-flush",
user_id="user-flush",
response_cost=0.05,
budget_reservation=reservation,
)
assert redis_cache.store[counter_key] == pytest.approx(0.35)
assert reservation["finalized"] is True