From 1c238b630fd289e146b532e309e4e2c81f093622 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Mon, 6 Apr 2026 17:25:48 -0700 Subject: [PATCH] fix(tests): update upsert tests to reflect new update-in-place behavior for existing budget_id --- .../test_upsert_budget_membership.py | 108 +++++------------- 1 file changed, 28 insertions(+), 80 deletions(-) diff --git a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py index a36dc7ff2e3..59312c5fc1c 100644 --- a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py +++ b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py @@ -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)