mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
a8fba14e10
The reset job moved from zeroing spend to an atomic decrement of the amount
it cleared, so every batched write now carries {"decrement": <cleared>}
instead of 0. Four tests still pinned 0 and had been failing since, which
also meant they no longer checked the amount at all. Assert the decrement
equals each row's own pre-reset spend, so a wrong amount fails the test.
1177 lines
45 KiB
Python
1177 lines
45 KiB
Python
import asyncio
|
|
from datetime import datetime, timedelta, timezone
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
|
|
from litellm.proxy._types import LiteLLM_BudgetTableFull
|
|
|
|
|
|
from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob
|
|
|
|
# Note: In our "fake" items we use dicts with fields that our fake reset functions modify.
|
|
# In a real-world scenario, these would be instances of LiteLLM_VerificationToken, LiteLLM_UserTable, etc.
|
|
|
|
|
|
def _attrify(d: dict):
|
|
"""
|
|
Wrap a dict so that attribute access (`.token`, `.user_id`, `.team_id`,
|
|
etc.) works alongside the existing item-access the fake_reset_* helpers
|
|
rely on. The reset job's narrow-write helpers use `getattr(item, "token",
|
|
None)` (et al), which returns None for plain dicts — that would silently
|
|
skip the row.
|
|
"""
|
|
|
|
class _AttrDict(dict):
|
|
def __getattr__(self, k):
|
|
try:
|
|
return self[k]
|
|
except KeyError:
|
|
raise AttributeError(k)
|
|
|
|
def __setattr__(self, k, v):
|
|
self[k] = v
|
|
|
|
return _AttrDict(d)
|
|
|
|
|
|
def _wire_batcher_for_test(prisma_client, fail_commit=False):
|
|
"""
|
|
Wire prisma_client.db.batch_() to return a mock batcher whose .commit() is
|
|
awaitable and whose per-table .update()/.update_many() calls get captured.
|
|
The reset job writes every reset through prisma.db.batch_() — key/user/team
|
|
rows one by one, and the budget tier's cascade as a single transaction — so
|
|
tests must let that batch path complete.
|
|
|
|
Only committed batches contribute to the returned list, mirroring prisma:
|
|
with fail_commit=True the transaction blows up and must persist nothing.
|
|
|
|
Returns the list that will accumulate {table, op, where, data} dicts from
|
|
each captured write.
|
|
"""
|
|
batch_calls = []
|
|
|
|
def make_batcher():
|
|
queued = []
|
|
|
|
class _Table:
|
|
def __init__(self, table_name):
|
|
self._table_name = table_name
|
|
|
|
def update(self, where=None, data=None):
|
|
queued.append(
|
|
{
|
|
"table": self._table_name,
|
|
"op": "update",
|
|
"where": where,
|
|
"data": data,
|
|
}
|
|
)
|
|
|
|
def update_many(self, where=None, data=None):
|
|
queued.append(
|
|
{
|
|
"table": self._table_name,
|
|
"op": "update_many",
|
|
"where": where,
|
|
"data": data,
|
|
}
|
|
)
|
|
|
|
async def commit():
|
|
if fail_commit:
|
|
raise RuntimeError("simulated Postgres failure committing the batch")
|
|
batch_calls.extend(queued)
|
|
|
|
batcher = MagicMock()
|
|
batcher.litellm_verificationtoken = _Table("key")
|
|
batcher.litellm_usertable = _Table("user")
|
|
batcher.litellm_teamtable = _Table("team")
|
|
batcher.litellm_budgettable = _Table("budget")
|
|
batcher.litellm_teammembership = _Table("team_membership")
|
|
batcher.litellm_organizationtable = _Table("org")
|
|
batcher.litellm_tagtable = _Table("tag")
|
|
batcher.litellm_endusertable = _Table("enduser")
|
|
batcher.commit = commit
|
|
return batcher
|
|
|
|
prisma_client.db.batch_ = MagicMock(side_effect=make_batcher)
|
|
return batch_calls
|
|
|
|
|
|
def _wire_cascade_reads_for_test(prisma_client):
|
|
"""
|
|
The budget tier's cascade reads the rows it is about to zero, so their
|
|
spend counters can be invalidated after the commit. Give each of those
|
|
tables an awaitable find_many so the reads resolve instead of falling into
|
|
the job's warn-and-continue path.
|
|
"""
|
|
for table in (
|
|
"litellm_teammembership",
|
|
"litellm_verificationtoken",
|
|
"litellm_organizationtable",
|
|
"litellm_tagtable",
|
|
"litellm_endusertable",
|
|
):
|
|
getattr(prisma_client.db, table).find_many = AsyncMock(return_value=[])
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reset_budget_keys_partial_failure():
|
|
"""
|
|
Test that if one key fails to reset, the failure for that key does not block processing of the other keys.
|
|
We simulate two keys where the first fails and the second succeeds.
|
|
"""
|
|
# Arrange
|
|
key1 = {
|
|
"id": "key1",
|
|
"spend": 10.0,
|
|
"budget_duration": 60,
|
|
} # Will trigger simulated failure
|
|
key2 = {"id": "key2", "spend": 15.0, "budget_duration": 60} # Should be updated
|
|
key3 = {"id": "key3", "spend": 20.0, "budget_duration": 60} # Should be updated
|
|
key4 = {"id": "key4", "spend": 25.0, "budget_duration": 60} # Should be updated
|
|
key5 = {"id": "key5", "spend": 30.0, "budget_duration": 60} # Should be updated
|
|
key6 = {"id": "key6", "spend": 35.0, "budget_duration": 60} # Should be updated
|
|
|
|
prisma_client = MagicMock()
|
|
prisma_client.get_data = AsyncMock(
|
|
return_value=[key1, key2, key3, key4, key5, key6]
|
|
)
|
|
prisma_client.update_data = AsyncMock()
|
|
# Reset job writes key resets via prisma.db.batch_().<table>.update — not
|
|
# via update_data — so wire that path.
|
|
batch_calls = _wire_batcher_for_test(prisma_client)
|
|
|
|
# Using a dummy logging object with async hooks mocked out.
|
|
proxy_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock()
|
|
|
|
job = ResetBudgetJob(proxy_logging_obj, prisma_client)
|
|
|
|
now = datetime.utcnow()
|
|
|
|
# token is needed because the new write path uses where={"token": ...}
|
|
# and _AttrDict makes getattr work alongside item access used by fake_reset_key.
|
|
for k in [key1, key2, key3, key4, key5, key6]:
|
|
k.setdefault("token", k["id"])
|
|
key1, key2, key3, key4, key5, key6 = (
|
|
_attrify(k) for k in [key1, key2, key3, key4, key5, key6]
|
|
)
|
|
pre_reset_spend = {
|
|
k["token"]: k["spend"] for k in [key2, key3, key4, key5, key6]
|
|
}
|
|
prisma_client.get_data = AsyncMock(
|
|
return_value=[key1, key2, key3, key4, key5, key6]
|
|
)
|
|
|
|
async def fake_reset_key(key, current_time, reset_settings=None):
|
|
if key["id"] == "key1":
|
|
# Simulate a failure on key1 (for example, this might be due to an invariant check)
|
|
raise Exception("Simulated failure for key1")
|
|
else:
|
|
# Simulate successful reset modification
|
|
key["spend"] = 0.0
|
|
# Compute a new reset time based on the budget duration
|
|
key["budget_reset_at"] = (
|
|
current_time + timedelta(seconds=key["budget_duration"])
|
|
).isoformat()
|
|
return key
|
|
|
|
with patch.object(
|
|
ResetBudgetJob, "_reset_budget_for_key", side_effect=fake_reset_key
|
|
) as mock_reset_key:
|
|
# Call the method; even though one key fails, the loop should process both
|
|
await job.reset_budget_for_litellm_keys()
|
|
# Allow any created tasks (logging hooks) to schedule
|
|
await asyncio.sleep(0.1)
|
|
|
|
# Assert that the helper was called for 6 keys
|
|
assert mock_reset_key.call_count == 6
|
|
|
|
# Assert that the new narrow write path got 5 batched updates (key1 failed).
|
|
# update_data must NOT have been called for keys.
|
|
prisma_client.update_data.assert_not_awaited()
|
|
key_writes = [c for c in batch_calls if c["table"] == "key"]
|
|
assert len(key_writes) == 5
|
|
written_ids = [c["where"]["token"] for c in key_writes]
|
|
assert written_ids == ["key2", "key3", "key4", "key5", "key6"]
|
|
# And every write must carry only {spend, budget_reset_at} — never the full row.
|
|
for c in key_writes:
|
|
assert set(c["data"].keys()) == {"spend", "budget_reset_at"}
|
|
assert c["data"]["spend"] == {"decrement": pre_reset_spend[c["where"]["token"]]}
|
|
|
|
# Verify that the failure logging hook was scheduled (due to the failure for key1)
|
|
failure_hook_calls = (
|
|
proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args_list
|
|
)
|
|
# There should be one failure hook call for keys (with call_type "reset_budget_keys")
|
|
assert any(
|
|
call.kwargs.get("call_type") == "reset_budget_keys"
|
|
for call in failure_hook_calls
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reset_budget_users_partial_failure():
|
|
"""
|
|
Test that if one user fails to reset, the reset loop still processes the other users.
|
|
We simulate two users where the first fails and the second is updated.
|
|
"""
|
|
user1 = {
|
|
"id": "user1",
|
|
"spend": 20.0,
|
|
"budget_duration": 120,
|
|
} # Will trigger simulated failure
|
|
user2 = {"id": "user2", "spend": 25.0, "budget_duration": 120} # Should be updated
|
|
user3 = {"id": "user3", "spend": 30.0, "budget_duration": 120} # Should be updated
|
|
user4 = {"id": "user4", "spend": 35.0, "budget_duration": 120} # Should be updated
|
|
user5 = {"id": "user5", "spend": 40.0, "budget_duration": 120} # Should be updated
|
|
user6 = {"id": "user6", "spend": 45.0, "budget_duration": 120} # Should be updated
|
|
|
|
prisma_client = MagicMock()
|
|
prisma_client.get_data = AsyncMock(
|
|
return_value=[user1, user2, user3, user4, user5, user6]
|
|
)
|
|
prisma_client.update_data = AsyncMock()
|
|
batch_calls = _wire_batcher_for_test(prisma_client)
|
|
|
|
proxy_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock()
|
|
|
|
job = ResetBudgetJob(proxy_logging_obj, prisma_client)
|
|
|
|
# user_id required for the new write path's where clause; _AttrDict so
|
|
# getattr(u, 'user_id') works alongside the dict access fake_reset_user uses.
|
|
for u in [user1, user2, user3, user4, user5, user6]:
|
|
u.setdefault("user_id", u["id"])
|
|
user1, user2, user3, user4, user5, user6 = (
|
|
_attrify(u) for u in [user1, user2, user3, user4, user5, user6]
|
|
)
|
|
pre_reset_spend = {
|
|
u["user_id"]: u["spend"] for u in [user2, user3, user4, user5, user6]
|
|
}
|
|
prisma_client.get_data = AsyncMock(
|
|
return_value=[user1, user2, user3, user4, user5, user6]
|
|
)
|
|
|
|
async def fake_reset_user(user, current_time, reset_settings=None):
|
|
if user["id"] == "user1":
|
|
raise Exception("Simulated failure for user1")
|
|
else:
|
|
user["spend"] = 0.0
|
|
user["budget_reset_at"] = (
|
|
current_time + timedelta(seconds=user["budget_duration"])
|
|
).isoformat()
|
|
return user
|
|
|
|
with patch.object(
|
|
ResetBudgetJob, "_reset_budget_for_user", side_effect=fake_reset_user
|
|
) as mock_reset_user:
|
|
await job.reset_budget_for_litellm_users()
|
|
await asyncio.sleep(0.1)
|
|
|
|
assert mock_reset_user.call_count == 6
|
|
prisma_client.update_data.assert_not_awaited()
|
|
user_writes = [c for c in batch_calls if c["table"] == "user"]
|
|
assert len(user_writes) == 5
|
|
written_ids = [c["where"]["user_id"] for c in user_writes]
|
|
assert written_ids == ["user2", "user3", "user4", "user5", "user6"]
|
|
for c in user_writes:
|
|
assert set(c["data"].keys()) == {"spend", "budget_reset_at"}
|
|
assert c["data"]["spend"] == {
|
|
"decrement": pre_reset_spend[c["where"]["user_id"]]
|
|
}
|
|
|
|
failure_hook_calls = (
|
|
proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args_list
|
|
)
|
|
assert any(
|
|
call.kwargs.get("call_type") == "reset_budget_users"
|
|
for call in failure_hook_calls
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reset_budget_endusers_cascade_failure_is_all_or_nothing():
|
|
"""
|
|
A failure anywhere in the budget-tier cascade must persist nothing, so the
|
|
tier stays due and the next scheduler tick retries it. Before the fix the
|
|
job committed the new budget_reset_at first and zeroed the dependent spend
|
|
afterwards, so a failure here left the tier stamped for the next window
|
|
while every end user stayed at the cap.
|
|
"""
|
|
endusers = [
|
|
_attrify({"user_id": f"user{i}", "spend": 20.0 + i, "budget_id": "budget1"})
|
|
for i in range(1, 7)
|
|
]
|
|
|
|
budget1 = LiteLLM_BudgetTableFull(
|
|
**{
|
|
"budget_id": "budget1",
|
|
"max_budget": 65.0,
|
|
"budget_duration": "2d",
|
|
"created_at": datetime.now(timezone.utc) - timedelta(days=3),
|
|
}
|
|
)
|
|
|
|
prisma_client = MagicMock()
|
|
|
|
async def get_data_mock(table_name, *args, **kwargs):
|
|
if table_name == "budget":
|
|
return [budget1]
|
|
elif table_name == "enduser":
|
|
return endusers
|
|
return []
|
|
|
|
prisma_client.get_data = AsyncMock()
|
|
prisma_client.get_data.side_effect = get_data_mock
|
|
prisma_client.update_data = AsyncMock()
|
|
batch_calls = _wire_batcher_for_test(prisma_client, fail_commit=True)
|
|
|
|
proxy_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock()
|
|
|
|
job = ResetBudgetJob(proxy_logging_obj, prisma_client)
|
|
|
|
await job.reset_budget_for_litellm_budget_table()
|
|
await asyncio.sleep(0.1)
|
|
|
|
assert batch_calls == [], "a failed cascade must not persist any write"
|
|
assert (
|
|
prisma_client.update_data.await_count == 0
|
|
), "budget_reset_at must not be advanced outside the cascade transaction"
|
|
|
|
failure_hook_calls = (
|
|
proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args_list
|
|
)
|
|
assert any(
|
|
call.kwargs.get("call_type") == "reset_budget_endusers"
|
|
for call in failure_hook_calls
|
|
)
|
|
proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reset_budget_endusers_are_zeroed_with_the_budget_window_advance():
|
|
"""
|
|
The happy path: every end user the tier gates is zeroed and the tier's
|
|
budget_reset_at advances, all inside one transaction.
|
|
"""
|
|
endusers = [
|
|
_attrify({"user_id": f"user{i}", "spend": 20.0 + i, "budget_id": "budget1"})
|
|
for i in range(1, 7)
|
|
]
|
|
|
|
budget1 = LiteLLM_BudgetTableFull(
|
|
**{
|
|
"budget_id": "budget1",
|
|
"max_budget": 65.0,
|
|
"budget_duration": "2d",
|
|
"created_at": datetime.now(timezone.utc) - timedelta(days=3),
|
|
}
|
|
)
|
|
|
|
prisma_client = MagicMock()
|
|
|
|
async def get_data_mock(table_name, *args, **kwargs):
|
|
if table_name == "budget":
|
|
return [budget1]
|
|
elif table_name == "enduser":
|
|
return endusers
|
|
return []
|
|
|
|
prisma_client.get_data = AsyncMock()
|
|
prisma_client.get_data.side_effect = get_data_mock
|
|
prisma_client.update_data = AsyncMock()
|
|
batch_calls = _wire_batcher_for_test(prisma_client)
|
|
|
|
proxy_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock()
|
|
|
|
job = ResetBudgetJob(proxy_logging_obj, prisma_client)
|
|
|
|
await job.reset_budget_for_litellm_budget_table()
|
|
await asyncio.sleep(0.1)
|
|
|
|
assert prisma_client.db.batch_.call_count == 1, "the cascade must be one transaction"
|
|
|
|
enduser_writes = [c for c in batch_calls if c["table"] == "enduser"]
|
|
assert len(enduser_writes) == 1
|
|
assert enduser_writes[0]["where"] == {"budget_id": {"in": ["budget1"]}, "spend": {"gt": 0}}
|
|
assert enduser_writes[0]["data"] == {"spend": 0}
|
|
|
|
budget_writes = [c for c in batch_calls if c["table"] == "budget"]
|
|
assert len(budget_writes) == 1
|
|
assert budget_writes[0]["where"] == {"budget_id": "budget1"}
|
|
assert budget_writes[0]["data"]["budget_reset_at"] > datetime.now(timezone.utc)
|
|
|
|
proxy_logging_obj.service_logging_obj.async_service_failure_hook.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reset_budget_teams_partial_failure():
|
|
"""
|
|
Test that if one team fails to reset, the loop processes both teams and only updates the ones that succeeded.
|
|
We simulate two teams where the first fails and the second is updated.
|
|
"""
|
|
team1 = {
|
|
"id": "team1",
|
|
"spend": 30.0,
|
|
"budget_duration": 180,
|
|
} # Will trigger simulated failure
|
|
team2 = {"id": "team2", "spend": 35.0, "budget_duration": 180} # Should be updated
|
|
|
|
prisma_client = MagicMock()
|
|
prisma_client.get_data = AsyncMock(return_value=[team1, team2])
|
|
prisma_client.update_data = AsyncMock()
|
|
batch_calls = _wire_batcher_for_test(prisma_client)
|
|
|
|
proxy_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock()
|
|
|
|
job = ResetBudgetJob(proxy_logging_obj, prisma_client)
|
|
|
|
# team_id required for the new write path's where clause; _AttrDict for getattr.
|
|
for t in [team1, team2]:
|
|
t.setdefault("team_id", t["id"])
|
|
team1, team2 = _attrify(team1), _attrify(team2)
|
|
pre_reset_spend = team2["spend"]
|
|
prisma_client.get_data = AsyncMock(return_value=[team1, team2])
|
|
|
|
async def fake_reset_team(team, current_time, reset_settings=None):
|
|
if team["id"] == "team1":
|
|
raise Exception("Simulated failure for team1")
|
|
else:
|
|
team["spend"] = 0.0
|
|
team["budget_reset_at"] = (
|
|
current_time + timedelta(seconds=team["budget_duration"])
|
|
).isoformat()
|
|
return team
|
|
|
|
with patch.object(
|
|
ResetBudgetJob, "_reset_budget_for_team", side_effect=fake_reset_team
|
|
) as mock_reset_team:
|
|
await job.reset_budget_for_litellm_teams()
|
|
await asyncio.sleep(0.1)
|
|
|
|
assert mock_reset_team.call_count == 2
|
|
prisma_client.update_data.assert_not_awaited()
|
|
team_writes = [c for c in batch_calls if c["table"] == "team"]
|
|
assert len(team_writes) == 1
|
|
assert team_writes[0]["where"] == {"team_id": "team2"}
|
|
assert set(team_writes[0]["data"].keys()) == {"spend", "budget_reset_at"}
|
|
assert team_writes[0]["data"]["spend"] == {"decrement": pre_reset_spend}
|
|
|
|
failure_hook_calls = (
|
|
proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args_list
|
|
)
|
|
assert any(
|
|
call.kwargs.get("call_type") == "reset_budget_teams"
|
|
for call in failure_hook_calls
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reset_budget_continues_other_categories_on_failure():
|
|
"""
|
|
Test that executing the overall reset_budget() method continues to process keys, users, and teams,
|
|
even if one of the sub-categories (here, users) experiences a partial failure.
|
|
|
|
In this simulation:
|
|
- All keys are processed successfully.
|
|
- One of the two users fails.
|
|
- All teams are processed successfully.
|
|
|
|
We then assert that:
|
|
- update_data is called for each category with the correctly updated items.
|
|
- Each get_data call is made (indicating that one failing category did not abort the others).
|
|
"""
|
|
# Arrange dummy items for each table
|
|
key1 = {"id": "key1", "spend": 10.0, "budget_duration": 60}
|
|
key2 = {"id": "key2", "spend": 15.0, "budget_duration": 60}
|
|
user1 = {
|
|
"id": "user1",
|
|
"spend": 20.0,
|
|
"budget_duration": 120,
|
|
} # Will fail in user reset
|
|
user2 = {"id": "user2", "spend": 25.0, "budget_duration": 120} # Succeeds
|
|
team1 = {"id": "team1", "spend": 30.0, "budget_duration": 180}
|
|
team2 = {"id": "team2", "spend": 35.0, "budget_duration": 180}
|
|
enduser1 = {"user_id": "user1", "spend": 25.0, "budget_id": "budget1"}
|
|
budget1 = LiteLLM_BudgetTableFull(
|
|
**{
|
|
"budget_id": "budget1",
|
|
"max_budget": 65.0,
|
|
"budget_duration": "2d",
|
|
"created_at": datetime.now(timezone.utc) - timedelta(days=3),
|
|
}
|
|
)
|
|
|
|
prisma_client = MagicMock()
|
|
|
|
async def fake_get_data(*, table_name, query_type, **kwargs):
|
|
if table_name == "key":
|
|
return [key1, key2]
|
|
elif table_name == "user":
|
|
return [user1, user2]
|
|
elif table_name == "team":
|
|
return [team1, team2]
|
|
elif table_name == "budget":
|
|
return [budget1]
|
|
elif table_name == "enduser":
|
|
return [enduser1]
|
|
return []
|
|
|
|
prisma_client.get_data = AsyncMock(side_effect=fake_get_data)
|
|
prisma_client.update_data = AsyncMock()
|
|
batch_calls = _wire_batcher_for_test(prisma_client)
|
|
# ID fields required by the new write path's where clauses; _AttrDict
|
|
# lets getattr() see them alongside the item-access fake_reset_* helpers use.
|
|
for k in [key1, key2]:
|
|
k.setdefault("token", k["id"])
|
|
for u in [user1, user2]:
|
|
u.setdefault("user_id", u["id"])
|
|
for t in [team1, team2]:
|
|
t.setdefault("team_id", t["id"])
|
|
key1, key2 = _attrify(key1), _attrify(key2)
|
|
user1, user2 = _attrify(user1), _attrify(user2)
|
|
team1, team2 = _attrify(team1), _attrify(team2)
|
|
enduser1 = _attrify(enduser1)
|
|
pre_reset_spend = {
|
|
**{k["token"]: k["spend"] for k in [key1, key2]},
|
|
**{u["user_id"]: u["spend"] for u in [user2]},
|
|
**{t["team_id"]: t["spend"] for t in [team1, team2]},
|
|
}
|
|
_wire_cascade_reads_for_test(prisma_client)
|
|
|
|
proxy_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock()
|
|
|
|
job = ResetBudgetJob(proxy_logging_obj, prisma_client)
|
|
|
|
async def fake_reset_key(key, current_time, reset_settings=None):
|
|
key["spend"] = 0.0
|
|
key["budget_reset_at"] = (
|
|
current_time + timedelta(seconds=key["budget_duration"])
|
|
).isoformat()
|
|
return key
|
|
|
|
async def fake_reset_user(user, current_time, reset_settings=None):
|
|
if user["id"] == "user1":
|
|
raise Exception("Simulated failure for user1")
|
|
user["spend"] = 0.0
|
|
user["budget_reset_at"] = (
|
|
current_time + timedelta(seconds=user["budget_duration"])
|
|
).isoformat()
|
|
return user
|
|
|
|
async def fake_reset_team(team, current_time, reset_settings=None):
|
|
team["spend"] = 0.0
|
|
team["budget_reset_at"] = (
|
|
current_time + timedelta(seconds=team["budget_duration"])
|
|
).isoformat()
|
|
return team
|
|
|
|
with (
|
|
patch.object(
|
|
ResetBudgetJob, "_reset_budget_for_key", side_effect=fake_reset_key
|
|
) as mock_reset_key,
|
|
patch.object(
|
|
ResetBudgetJob, "_reset_budget_for_user", side_effect=fake_reset_user
|
|
) as mock_reset_user,
|
|
patch.object(
|
|
ResetBudgetJob, "_reset_budget_for_team", side_effect=fake_reset_team
|
|
) as mock_reset_team,
|
|
):
|
|
# Call the overall reset_budget method.
|
|
await job.reset_budget()
|
|
await asyncio.sleep(0.1)
|
|
|
|
# Verify that get_data was called for each table. We can check the table names across calls.
|
|
called_tables = {
|
|
call.kwargs.get("table_name") for call in prisma_client.get_data.await_args_list
|
|
}
|
|
assert called_tables == {"key", "user", "team", "budget", "enduser"}
|
|
|
|
# Every category writes through the batch path now, so update_data is unused.
|
|
prisma_client.update_data.assert_not_awaited()
|
|
|
|
# The budget tier's cascade still ran despite the failing user category.
|
|
assert len([c for c in batch_calls if c["table"] == "team_membership"]) == 1
|
|
enduser_writes = [c for c in batch_calls if c["table"] == "enduser"]
|
|
assert len(enduser_writes) == 1
|
|
assert enduser_writes[0]["where"] == {"budget_id": {"in": ["budget1"]}, "spend": {"gt": 0}}
|
|
assert enduser_writes[0]["data"] == {"spend": 0}
|
|
|
|
# Check the new batch write path: 2 keys + 1 user (user1 failed) + 2 teams.
|
|
# `op` separates the per-row resets from the cascade sweep, which also
|
|
# targets the key table.
|
|
key_writes = [c for c in batch_calls if c["table"] == "key" and c["op"] == "update"]
|
|
user_writes = [c for c in batch_calls if c["table"] == "user"]
|
|
team_writes = [c for c in batch_calls if c["table"] == "team"]
|
|
assert len(key_writes) == 2
|
|
assert len(user_writes) == 1
|
|
assert user_writes[0]["where"] == {"user_id": "user2"}
|
|
assert len(team_writes) == 2
|
|
# Every batched write must carry only the two reset fields, never the full row.
|
|
for c in key_writes + user_writes + team_writes:
|
|
assert set(c["data"].keys()) == {"spend", "budget_reset_at"}
|
|
assert c["data"]["spend"] == {
|
|
"decrement": pre_reset_spend[next(iter(c["where"].values()))]
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Additional tests for service logger behavior (keys, users, teams, endusers)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_service_logger_keys_success():
|
|
"""
|
|
Test that when resetting keys succeeds (all keys are updated) the service
|
|
logger success hook is called with the correct event metadata and no exception is logged.
|
|
"""
|
|
keys = [
|
|
_attrify(
|
|
{"id": "key1", "spend": 10.0, "budget_duration": 60, "token": "key1"}
|
|
),
|
|
_attrify(
|
|
{"id": "key2", "spend": 15.0, "budget_duration": 60, "token": "key2"}
|
|
),
|
|
]
|
|
prisma_client = MagicMock()
|
|
prisma_client.get_data = AsyncMock(return_value=keys)
|
|
prisma_client.update_data = AsyncMock()
|
|
_wire_batcher_for_test(prisma_client)
|
|
|
|
proxy_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock()
|
|
|
|
job = ResetBudgetJob(proxy_logging_obj, prisma_client)
|
|
|
|
async def fake_reset_key(key, current_time, reset_settings=None):
|
|
key["spend"] = 0.0
|
|
key["budget_reset_at"] = (
|
|
current_time + timedelta(seconds=key["budget_duration"])
|
|
).isoformat()
|
|
return key
|
|
|
|
with patch.object(
|
|
ResetBudgetJob,
|
|
"_reset_budget_for_key",
|
|
side_effect=fake_reset_key,
|
|
):
|
|
with patch(
|
|
"litellm.proxy.common_utils.reset_budget_job.verbose_proxy_logger.exception"
|
|
) as mock_verbose_exc:
|
|
await job.reset_budget_for_litellm_keys()
|
|
# Allow async logging task to complete
|
|
await asyncio.sleep(0.1)
|
|
mock_verbose_exc.assert_not_called()
|
|
|
|
# Verify success hook call
|
|
proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_called_once()
|
|
(
|
|
args,
|
|
kwargs,
|
|
) = proxy_logging_obj.service_logging_obj.async_service_success_hook.call_args
|
|
event_metadata = kwargs.get("event_metadata", {})
|
|
assert event_metadata.get("num_keys_found") == len(keys)
|
|
assert event_metadata.get("num_keys_updated") == len(keys)
|
|
assert event_metadata.get("num_keys_failed") == 0
|
|
# Failure hook should not be executed.
|
|
proxy_logging_obj.service_logging_obj.async_service_failure_hook.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_service_logger_keys_failure():
|
|
"""
|
|
Test that when a key reset fails the service logger failure hook is called,
|
|
the event metadata reflects the number of keys processed, and that the verbose
|
|
logger exception is called.
|
|
"""
|
|
keys = [
|
|
{"id": "key1", "spend": 10.0, "budget_duration": 60},
|
|
{"id": "key2", "spend": 15.0, "budget_duration": 60},
|
|
]
|
|
prisma_client = MagicMock()
|
|
prisma_client.get_data = AsyncMock(return_value=keys)
|
|
prisma_client.update_data = AsyncMock()
|
|
|
|
proxy_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock()
|
|
|
|
job = ResetBudgetJob(proxy_logging_obj, prisma_client)
|
|
|
|
async def fake_reset_key(key, current_time, reset_settings=None):
|
|
if key["id"] == "key1":
|
|
raise Exception("Simulated failure for key1")
|
|
key["spend"] = 0.0
|
|
key["budget_reset_at"] = (
|
|
current_time + timedelta(seconds=key["budget_duration"])
|
|
).isoformat()
|
|
return key
|
|
|
|
with patch.object(
|
|
ResetBudgetJob,
|
|
"_reset_budget_for_key",
|
|
side_effect=fake_reset_key,
|
|
):
|
|
with patch(
|
|
"litellm.proxy.common_utils.reset_budget_job.verbose_proxy_logger.exception"
|
|
) as mock_verbose_exc:
|
|
await job.reset_budget_for_litellm_keys()
|
|
await asyncio.sleep(0.1)
|
|
# Expect at least one exception logged (the inner error and the outer catch)
|
|
assert mock_verbose_exc.call_count >= 1
|
|
# Verify exception was logged with correct message
|
|
assert any(
|
|
"Failed to reset budget for key" in str(call.args)
|
|
for call in mock_verbose_exc.call_args_list
|
|
)
|
|
|
|
proxy_logging_obj.service_logging_obj.async_service_failure_hook.assert_called_once()
|
|
(
|
|
args,
|
|
kwargs,
|
|
) = proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args
|
|
event_metadata = kwargs.get("event_metadata", {})
|
|
assert event_metadata.get("num_keys_found") == len(keys)
|
|
# the row payload is deliberately absent: serializing every found row on the
|
|
# event loop is what blocked auth on the sweeping pod
|
|
assert "keys_found" not in event_metadata
|
|
# Success hook should not be called.
|
|
proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_service_logger_users_success():
|
|
"""
|
|
Test that when resetting users succeeds the service logger success hook is called with
|
|
the correct metadata and no exception is logged.
|
|
"""
|
|
users = [
|
|
_attrify(
|
|
{"id": "user1", "spend": 20.0, "budget_duration": 120, "user_id": "user1"}
|
|
),
|
|
_attrify(
|
|
{"id": "user2", "spend": 25.0, "budget_duration": 120, "user_id": "user2"}
|
|
),
|
|
]
|
|
prisma_client = MagicMock()
|
|
prisma_client.get_data = AsyncMock(return_value=users)
|
|
prisma_client.update_data = AsyncMock()
|
|
_wire_batcher_for_test(prisma_client)
|
|
|
|
proxy_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock()
|
|
|
|
job = ResetBudgetJob(proxy_logging_obj, prisma_client)
|
|
|
|
async def fake_reset_user(user, current_time, reset_settings=None):
|
|
user["spend"] = 0.0
|
|
user["budget_reset_at"] = (
|
|
current_time + timedelta(seconds=user["budget_duration"])
|
|
).isoformat()
|
|
return user
|
|
|
|
with patch.object(
|
|
ResetBudgetJob,
|
|
"_reset_budget_for_user",
|
|
side_effect=fake_reset_user,
|
|
):
|
|
with patch(
|
|
"litellm.proxy.common_utils.reset_budget_job.verbose_proxy_logger.exception"
|
|
) as mock_verbose_exc:
|
|
await job.reset_budget_for_litellm_users()
|
|
await asyncio.sleep(0.1)
|
|
mock_verbose_exc.assert_not_called()
|
|
|
|
proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_called_once()
|
|
(
|
|
args,
|
|
kwargs,
|
|
) = proxy_logging_obj.service_logging_obj.async_service_success_hook.call_args
|
|
event_metadata = kwargs.get("event_metadata", {})
|
|
assert event_metadata.get("num_users_found") == len(users)
|
|
assert event_metadata.get("num_users_updated") == len(users)
|
|
assert event_metadata.get("num_users_failed") == 0
|
|
proxy_logging_obj.service_logging_obj.async_service_failure_hook.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_service_logger_users_failure():
|
|
"""
|
|
Test that a failure during user reset calls the failure hook with appropriate metadata,
|
|
logs the exception, and does not call the success hook.
|
|
"""
|
|
users = [
|
|
{"id": "user1", "spend": 20.0, "budget_duration": 120},
|
|
{"id": "user2", "spend": 25.0, "budget_duration": 120},
|
|
]
|
|
prisma_client = MagicMock()
|
|
prisma_client.get_data = AsyncMock(return_value=users)
|
|
prisma_client.update_data = AsyncMock()
|
|
|
|
proxy_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock()
|
|
|
|
job = ResetBudgetJob(proxy_logging_obj, prisma_client)
|
|
|
|
async def fake_reset_user(user, current_time, reset_settings=None):
|
|
if user["id"] == "user1":
|
|
raise Exception("Simulated failure for user1")
|
|
user["spend"] = 0.0
|
|
user["budget_reset_at"] = (
|
|
current_time + timedelta(seconds=user["budget_duration"])
|
|
).isoformat()
|
|
return user
|
|
|
|
with patch.object(
|
|
ResetBudgetJob,
|
|
"_reset_budget_for_user",
|
|
side_effect=fake_reset_user,
|
|
):
|
|
with patch(
|
|
"litellm.proxy.common_utils.reset_budget_job.verbose_proxy_logger.exception"
|
|
) as mock_verbose_exc:
|
|
await job.reset_budget_for_litellm_users()
|
|
await asyncio.sleep(0.1)
|
|
# Verify exception logging
|
|
assert mock_verbose_exc.call_count >= 1
|
|
# Verify exception was logged with correct message
|
|
assert any(
|
|
"Failed to reset budget for user" in str(call.args)
|
|
for call in mock_verbose_exc.call_args_list
|
|
)
|
|
|
|
proxy_logging_obj.service_logging_obj.async_service_failure_hook.assert_called_once()
|
|
(
|
|
args,
|
|
kwargs,
|
|
) = proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args
|
|
event_metadata = kwargs.get("event_metadata", {})
|
|
assert event_metadata.get("num_users_found") == len(users)
|
|
assert "users_found" not in event_metadata
|
|
proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_service_logger_teams_success():
|
|
"""
|
|
Test that when resetting teams is successful the service logger success hook is called with
|
|
the proper metadata and nothing is logged as an exception.
|
|
"""
|
|
teams = [
|
|
_attrify(
|
|
{"id": "team1", "spend": 30.0, "budget_duration": 180, "team_id": "team1"}
|
|
),
|
|
_attrify(
|
|
{"id": "team2", "spend": 35.0, "budget_duration": 180, "team_id": "team2"}
|
|
),
|
|
]
|
|
prisma_client = MagicMock()
|
|
prisma_client.get_data = AsyncMock(return_value=teams)
|
|
prisma_client.update_data = AsyncMock()
|
|
_wire_batcher_for_test(prisma_client)
|
|
|
|
proxy_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock()
|
|
|
|
job = ResetBudgetJob(proxy_logging_obj, prisma_client)
|
|
|
|
async def fake_reset_team(team, current_time, reset_settings=None):
|
|
team["spend"] = 0.0
|
|
team["budget_reset_at"] = (
|
|
current_time + timedelta(seconds=team["budget_duration"])
|
|
).isoformat()
|
|
return team
|
|
|
|
with patch.object(
|
|
ResetBudgetJob,
|
|
"_reset_budget_for_team",
|
|
side_effect=fake_reset_team,
|
|
):
|
|
with patch(
|
|
"litellm.proxy.common_utils.reset_budget_job.verbose_proxy_logger.exception"
|
|
) as mock_verbose_exc:
|
|
await job.reset_budget_for_litellm_teams()
|
|
await asyncio.sleep(0.1)
|
|
mock_verbose_exc.assert_not_called()
|
|
|
|
proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_called_once()
|
|
(
|
|
args,
|
|
kwargs,
|
|
) = proxy_logging_obj.service_logging_obj.async_service_success_hook.call_args
|
|
event_metadata = kwargs.get("event_metadata", {})
|
|
assert event_metadata.get("num_teams_found") == len(teams)
|
|
assert event_metadata.get("num_teams_updated") == len(teams)
|
|
assert event_metadata.get("num_teams_failed") == 0
|
|
proxy_logging_obj.service_logging_obj.async_service_failure_hook.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_service_logger_teams_failure():
|
|
"""
|
|
Test that a failure during team reset triggers the failure hook with proper metadata,
|
|
results in an exception log and no success hook call.
|
|
"""
|
|
teams = [
|
|
{"id": "team1", "spend": 30.0, "budget_duration": 180},
|
|
{"id": "team2", "spend": 35.0, "budget_duration": 180},
|
|
]
|
|
prisma_client = MagicMock()
|
|
prisma_client.get_data = AsyncMock(return_value=teams)
|
|
prisma_client.update_data = AsyncMock()
|
|
|
|
proxy_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock()
|
|
|
|
job = ResetBudgetJob(proxy_logging_obj, prisma_client)
|
|
|
|
async def fake_reset_team(team, current_time, reset_settings=None):
|
|
if team["id"] == "team1":
|
|
raise Exception("Simulated failure for team1")
|
|
team["spend"] = 0.0
|
|
team["budget_reset_at"] = (
|
|
current_time + timedelta(seconds=team["budget_duration"])
|
|
).isoformat()
|
|
return team
|
|
|
|
with patch.object(
|
|
ResetBudgetJob,
|
|
"_reset_budget_for_team",
|
|
side_effect=fake_reset_team,
|
|
):
|
|
with patch(
|
|
"litellm.proxy.common_utils.reset_budget_job.verbose_proxy_logger.exception"
|
|
) as mock_verbose_exc:
|
|
await job.reset_budget_for_litellm_teams()
|
|
await asyncio.sleep(0.1)
|
|
# Verify exception logging
|
|
assert mock_verbose_exc.call_count >= 1
|
|
# Verify exception was logged with correct message
|
|
assert any(
|
|
"Failed to reset budget for team" in str(call.args)
|
|
for call in mock_verbose_exc.call_args_list
|
|
)
|
|
|
|
proxy_logging_obj.service_logging_obj.async_service_failure_hook.assert_called_once()
|
|
(
|
|
args,
|
|
kwargs,
|
|
) = proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args
|
|
event_metadata = kwargs.get("event_metadata", {})
|
|
assert event_metadata.get("num_teams_found") == len(teams)
|
|
assert "teams_found" not in event_metadata
|
|
proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_service_logger_endusers_success():
|
|
"""
|
|
Test that when the budget-tier cascade commits, the service logger success
|
|
hook is called with the correct metadata and no exception is logged.
|
|
"""
|
|
endusers = [
|
|
_attrify({"user_id": "user1", "spend": 25.0, "budget_id": "budget1"}),
|
|
_attrify({"user_id": "user2", "spend": 25.0, "budget_id": "budget1"}),
|
|
]
|
|
budgets = [
|
|
LiteLLM_BudgetTableFull(
|
|
**{
|
|
"budget_id": "budget1",
|
|
"max_budget": 65.0,
|
|
"budget_duration": "2d",
|
|
"created_at": datetime.now(timezone.utc) - timedelta(days=3),
|
|
}
|
|
)
|
|
]
|
|
|
|
async def fake_get_data(*, table_name, query_type, **kwargs):
|
|
if table_name == "budget":
|
|
return budgets
|
|
elif table_name == "enduser":
|
|
return endusers
|
|
return []
|
|
|
|
prisma_client = MagicMock()
|
|
prisma_client.get_data = AsyncMock(side_effect=fake_get_data)
|
|
prisma_client.update_data = AsyncMock()
|
|
batch_calls = _wire_batcher_for_test(prisma_client)
|
|
_wire_cascade_reads_for_test(prisma_client)
|
|
|
|
proxy_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock()
|
|
|
|
job = ResetBudgetJob(proxy_logging_obj, prisma_client)
|
|
|
|
with patch(
|
|
"litellm.proxy.common_utils.reset_budget_job.verbose_proxy_logger.exception"
|
|
) as mock_verbose_exc:
|
|
await job.reset_budget_for_litellm_budget_table()
|
|
await asyncio.sleep(0.1)
|
|
mock_verbose_exc.assert_not_called()
|
|
|
|
enduser_writes = [c for c in batch_calls if c["table"] == "enduser"]
|
|
assert len(enduser_writes) == 1
|
|
assert enduser_writes[0]["where"] == {"budget_id": {"in": ["budget1"]}, "spend": {"gt": 0}}
|
|
|
|
proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_called_once()
|
|
(
|
|
args,
|
|
kwargs,
|
|
) = proxy_logging_obj.service_logging_obj.async_service_success_hook.call_args
|
|
event_metadata = kwargs.get("event_metadata", {})
|
|
assert event_metadata.get("num_budgets_found") == len(budgets)
|
|
assert event_metadata.get("num_endusers_found") == len(endusers)
|
|
assert event_metadata.get("num_endusers_updated") == len(endusers)
|
|
assert event_metadata.get("num_endusers_failed") == 0
|
|
proxy_logging_obj.service_logging_obj.async_service_failure_hook.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_service_logger_endusers_failure():
|
|
"""
|
|
Test that a failed cascade calls the failure hook with the rows it had
|
|
found, logs the exception, and does not call the success hook.
|
|
"""
|
|
endusers = [
|
|
_attrify({"user_id": "user1", "spend": 25.0, "budget_id": "budget1"}),
|
|
_attrify({"user_id": "user2", "spend": 25.0, "budget_id": "budget1"}),
|
|
]
|
|
budgets = [
|
|
LiteLLM_BudgetTableFull(
|
|
**{
|
|
"budget_id": "budget1",
|
|
"max_budget": 65.0,
|
|
"budget_duration": "2d",
|
|
"created_at": datetime.now(timezone.utc) - timedelta(days=3),
|
|
}
|
|
)
|
|
]
|
|
|
|
async def fake_get_data(*, table_name, query_type, **kwargs):
|
|
if table_name == "budget":
|
|
return budgets
|
|
elif table_name == "enduser":
|
|
return endusers
|
|
return []
|
|
|
|
prisma_client = MagicMock()
|
|
prisma_client.get_data = AsyncMock(side_effect=fake_get_data)
|
|
prisma_client.update_data = AsyncMock()
|
|
_wire_batcher_for_test(prisma_client, fail_commit=True)
|
|
_wire_cascade_reads_for_test(prisma_client)
|
|
|
|
proxy_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock()
|
|
|
|
job = ResetBudgetJob(proxy_logging_obj, prisma_client)
|
|
|
|
with patch(
|
|
"litellm.proxy.common_utils.reset_budget_job.verbose_proxy_logger.exception"
|
|
) as mock_verbose_exc:
|
|
await job.reset_budget_for_litellm_budget_table()
|
|
await asyncio.sleep(0.1)
|
|
# The log must name the whole cascade, not just end users: the write
|
|
# that failed could have been any of team member / enduser / org / tag
|
|
# spend or the budget_reset_at advance.
|
|
assert mock_verbose_exc.call_count == 1
|
|
assert "budget table cascade" in str(mock_verbose_exc.call_args.args[0])
|
|
|
|
proxy_logging_obj.service_logging_obj.async_service_failure_hook.assert_called_once()
|
|
(
|
|
args,
|
|
kwargs,
|
|
) = proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args
|
|
event_metadata = kwargs.get("event_metadata", {})
|
|
assert event_metadata.get("num_budgets_found") == len(budgets)
|
|
assert event_metadata.get("num_endusers_found") == len(endusers)
|
|
assert "endusers_found" not in event_metadata
|
|
assert "budgets_found" not in event_metadata
|
|
proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reset_budget_for_litellm_team_members_called():
|
|
"""
|
|
Test that when reset_budget_for_litellm_budget_table is called, team
|
|
members' spend is zeroed as part of the cascade transaction.
|
|
"""
|
|
# Arrange
|
|
budget1 = LiteLLM_BudgetTableFull(
|
|
**{
|
|
"budget_id": "budget1",
|
|
"max_budget": 100.0,
|
|
"budget_duration": "1d",
|
|
"created_at": datetime.now(timezone.utc) - timedelta(days=2),
|
|
}
|
|
)
|
|
|
|
enduser1 = _attrify({"user_id": "user1", "spend": 25.0, "budget_id": "budget1"})
|
|
|
|
prisma_client = MagicMock()
|
|
|
|
async def fake_get_data(*, table_name, query_type, **kwargs):
|
|
if table_name == "budget":
|
|
return [budget1]
|
|
elif table_name == "enduser":
|
|
return [enduser1]
|
|
return []
|
|
|
|
prisma_client.get_data = AsyncMock(side_effect=fake_get_data)
|
|
prisma_client.update_data = AsyncMock()
|
|
prisma_client.db = MagicMock()
|
|
batch_calls = _wire_batcher_for_test(prisma_client)
|
|
_wire_cascade_reads_for_test(prisma_client)
|
|
|
|
proxy_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj = MagicMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock()
|
|
proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock()
|
|
|
|
job = ResetBudgetJob(proxy_logging_obj, prisma_client)
|
|
|
|
# Act
|
|
await job.reset_budget_for_litellm_budget_table()
|
|
|
|
# Assert
|
|
team_member_writes = [c for c in batch_calls if c["table"] == "team_membership"]
|
|
assert len(team_member_writes) == 1
|
|
assert team_member_writes[0]["where"]["budget_id"]["in"] == ["budget1"]
|
|
assert team_member_writes[0]["data"] == {"spend": 0}
|