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

fix(tests): update upsert tests to reflect new update-in-place behavior for existing budget_id

Browse source
This commit is contained in:
Ishaan Jaffer 2026-04-06 17:25:48 -07:00
parent 9e0fb6bd17
commit 1c238b630f
No known key found for this signature in database
1 changed files with 28 additions and 80 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

108
tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py
View file

@ -63,50 +63,33 @@ async def test_upsert_disconnect(mock_tx, fake_user):
mock_tx.litellm_teammembership.upsert.assert_not_called()
# TEST: existing budget id, creates new budget (current behavior)
# TEST: existing budget id → updates budget in-place (current behavior)
@pytest.mark.asyncio
async def test_upsert_with_existing_budget_id_creates_new(mock_tx, fake_user):
"""
Test that even when existing_budget_id is provided, the function creates a new budget.
This reflects the current implementation behavior.
Test that when existing_budget_id is provided, the function updates the budget in-place.
"""
await _upsert_budget_and_membership(
mock_tx,
team_id="team-2",
user_id="user-2",
max_budget=42.0,
existing_budget_id="bud-999", # This parameter is currently unused
existing_budget_id="bud-999",
user_api_key_dict=fake_user,
)
# Should create a new budget, not update existing
mock_tx.litellm_budgettable.create.assert_awaited_once_with(
# Should update the existing budget, not create a new one
mock_tx.litellm_budgettable.update.assert_awaited_once_with(
where={"budget_id": "bud-999"},
data={
"max_budget": 42.0,
"created_by": fake_user.user_id,
"updated_by": fake_user.user_id,
},
include={"team_membership": True},
)
# Should upsert team membership with the new budget ID
new_budget_id = mock_tx.litellm_budgettable.create.return_value.budget_id
mock_tx.litellm_teammembership.upsert.assert_awaited_once_with(
where={"user_id_team_id": {"user_id": "user-2", "team_id": "team-2"}},
data={
"create": {
"user_id": "user-2",
"team_id": "team-2",
"litellm_budget_table": {"connect": {"budget_id": new_budget_id}},
},
"update": {
"litellm_budget_table": {"connect": {"budget_id": new_budget_id}},
},
},
)
# Should NOT update existing budget
mock_tx.litellm_budgettable.update.assert_not_called()
# Should NOT create a new budget or touch membership
mock_tx.litellm_budgettable.create.assert_not_called()
mock_tx.litellm_teammembership.upsert.assert_not_called()
mock_tx.litellm_teammembership.update.assert_not_called()
@ -176,62 +159,43 @@ async def test_upsert_create_then_create_another(mock_tx, fake_user):
mock_tx.litellm_budgettable.create.assert_awaited_once()
mock_tx.litellm_teammembership.upsert.assert_awaited_once()
# SECOND CALL – reset call history and create another budget
# SECOND CALL – reset call history; this time we supply the existing budget_id
mock_tx.litellm_budgettable.create.reset_mock()
mock_tx.litellm_teammembership.upsert.reset_mock()
mock_tx.litellm_budgettable.update.reset_mock()
# Set up a new budget ID for the second create call
mock_tx.litellm_budgettable.create.return_value = types.SimpleNamespace(budget_id="new-budget-456")
await _upsert_budget_and_membership(
mock_tx,
team_id="team-42",
user_id="user-42",
max_budget=25.0, # new limit
existing_budget_id=created_bid, # this is ignored in current implementation
max_budget=25.0,
existing_budget_id=created_bid, # now used: triggers in-place update
user_api_key_dict=fake_user,
)
# Should create another new budget (not update existing)
mock_tx.litellm_budgettable.create.assert_awaited_once_with(
# Should update the existing budget in-place, not create a new one
mock_tx.litellm_budgettable.update.assert_awaited_once_with(
where={"budget_id": created_bid},
data={
"max_budget": 25.0,
"created_by": fake_user.user_id,
"updated_by": fake_user.user_id,
},
include={"team_membership": True},
)
# Should upsert team membership with the new budget ID
new_budget_id = mock_tx.litellm_budgettable.create.return_value.budget_id
mock_tx.litellm_teammembership.upsert.assert_awaited_once_with(
where={"user_id_team_id": {"user_id": "user-42", "team_id": "team-42"}},
data={
"create": {
"user_id": "user-42",
"team_id": "team-42",
"litellm_budget_table": {"connect": {"budget_id": new_budget_id}},
},
"update": {
"litellm_budget_table": {"connect": {"budget_id": new_budget_id}},
},
},
)
# Should NOT call update
mock_tx.litellm_budgettable.update.assert_not_called()
# Should NOT create a new budget or touch membership
mock_tx.litellm_budgettable.create.assert_not_called()
mock_tx.litellm_teammembership.upsert.assert_not_called()
# TEST: update rpm_limit for member with existing budget_id
# TEST: update rpm_limit for member with existing budget_id → updates in-place
@pytest.mark.asyncio
async def test_upsert_rpm_limit_update_creates_new_budget(mock_tx, fake_user):
"""
Test that updating rpm_limit for a member with an existing budget_id
creates a new budget with the new rpm/tpm limits and assigns it to the user.
updates the existing budget in-place (not creates a new one).
"""
existing_budget_id = "existing-budget-456"
await _upsert_budget_and_membership(
mock_tx,
team_id="team-rpm-test",
@ -240,39 +204,23 @@ async def test_upsert_rpm_limit_update_creates_new_budget(mock_tx, fake_user):
existing_budget_id=existing_budget_id,
user_api_key_dict=fake_user,
tpm_limit=1000,
rpm_limit=100, # updating rpm_limit
rpm_limit=100,
)
# Should create a new budget with all the specified limits
mock_tx.litellm_budgettable.create.assert_awaited_once_with(
# Should update the existing budget with all specified limits
mock_tx.litellm_budgettable.update.assert_awaited_once_with(
where={"budget_id": existing_budget_id},
data={
"max_budget": 50.0,
"tpm_limit": 1000,
"rpm_limit": 100,
"created_by": fake_user.user_id,
"updated_by": fake_user.user_id,
},
include={"team_membership": True},
)
# Should NOT update the existing budget
mock_tx.litellm_budgettable.update.assert_not_called()
# Should upsert team membership with the new budget ID
new_budget_id = mock_tx.litellm_budgettable.create.return_value.budget_id
mock_tx.litellm_teammembership.upsert.assert_awaited_once_with(
where={"user_id_team_id": {"user_id": "user-rpm-test", "team_id": "team-rpm-test"}},
data={
"create": {
"user_id": "user-rpm-test",
"team_id": "team-rpm-test",
"litellm_budget_table": {"connect": {"budget_id": new_budget_id}},
},
"update": {
"litellm_budget_table": {"connect": {"budget_id": new_budget_id}},
},
},
)
# Should NOT create a new budget or touch membership
mock_tx.litellm_budgettable.create.assert_not_called()
mock_tx.litellm_teammembership.upsert.assert_not_called()
# TEST: create new budget with only rpm_limit (no max_budget)