From 512999f1249b00a02a30f049a0cfa36e829ff989 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 20 Jun 2025 16:20:04 -0700 Subject: [PATCH] test: add unit tests for default team member budget --- .../test_management_helpers_utils.py | 293 ++++++++++++++++++ 1 file changed, 293 insertions(+) create mode 100644 tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py diff --git a/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py b/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py new file mode 100644 index 00000000000..2f6ecb776db --- /dev/null +++ b/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py @@ -0,0 +1,293 @@ +import json +import os +import sys +import uuid +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../../..") +) # Adds the parent directory to the system path + +from litellm.proxy._types import ( + LiteLLM_TeamMembership, + LiteLLM_UserTable, + Member, + UserAPIKeyAuth, +) +from litellm.proxy.management_helpers.utils import add_new_member + + +@pytest.mark.asyncio +async def test_add_new_member_uses_default_team_budget_id(): + """ + Test that add_new_member uses the default_team_budget_id when max_budget_in_team is None. + + This test verifies that: + 1. When max_budget_in_team is None + 2. And default_team_budget_id is provided + 3. The team membership is created with the default_team_budget_id + """ + from litellm.proxy._types import LitellmUserRoles + + # Setup test data + test_user_id = "test_user_123" + test_team_id = "test_team_456" + test_default_budget_id = "default_budget_789" + test_admin_name = "test_admin" + + # Create a Member object with user_id + new_member = Member(user_id=test_user_id, role="user") + + # Create UserAPIKeyAuth object + user_api_key_dict = UserAPIKeyAuth( + user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + # Mock the prisma client + mock_prisma_client = AsyncMock() + + # Mock the user table upsert operation + mock_user_response = MagicMock() + mock_user_response.model_dump.return_value = { + "user_id": test_user_id, + "user_email": None, + "teams": [test_team_id], + "user_role": "internal_user", + } + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( + return_value=mock_user_response + ) + + # Mock the team membership creation + mock_team_membership_response = MagicMock() + mock_team_membership_response.model_dump.return_value = { + "team_id": test_team_id, + "user_id": test_user_id, + "budget_id": test_default_budget_id, + "litellm_budget_table": None, + } + mock_prisma_client.db.litellm_teammembership.create = AsyncMock( + return_value=mock_team_membership_response + ) + + # Call the function with max_budget_in_team=None and a default_team_budget_id + result_user, result_team_membership = await add_new_member( + new_member=new_member, + max_budget_in_team=None, # This is the key - no max budget specified + prisma_client=mock_prisma_client, + team_id=test_team_id, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=test_admin_name, + default_team_budget_id=test_default_budget_id, # This should be used + ) + + # Verify that the user was created/updated correctly + assert result_user is not None + assert result_user.user_id == test_user_id + + # Verify that the team membership was created correctly + assert result_team_membership is not None + assert result_team_membership.team_id == test_team_id + assert result_team_membership.user_id == test_user_id + assert result_team_membership.budget_id == test_default_budget_id + + # Verify that the prisma client methods were called correctly + mock_prisma_client.db.litellm_usertable.upsert.assert_called_once() + mock_prisma_client.db.litellm_teammembership.create.assert_called_once() + + # Verify that no budget table creation was called (since max_budget_in_team is None) + assert ( + not hasattr(mock_prisma_client.db, "litellm_budgettable") + or not mock_prisma_client.db.litellm_budgettable.create.called + ) + + # Verify the team membership was created with the correct budget_id + team_membership_call_args = ( + mock_prisma_client.db.litellm_teammembership.create.call_args + ) + assert team_membership_call_args is not None + create_data = team_membership_call_args.kwargs["data"] + assert create_data["budget_id"] == test_default_budget_id + + +@pytest.mark.asyncio +async def test_add_new_member_creates_new_budget_when_max_budget_provided(): + """ + Test that add_new_member creates a new budget when max_budget_in_team is provided. + + This test verifies that: + 1. When max_budget_in_team is provided + 2. A new budget is created in the litellm_budgettable + 3. The new budget_id is used for the team membership + """ + from litellm.proxy._types import LitellmUserRoles + + # Setup test data + test_user_id = "test_user_123" + test_team_id = "test_team_456" + test_max_budget = 100.0 + test_new_budget_id = "new_budget_789" + test_admin_name = "test_admin" + + # Create a Member object with user_id + new_member = Member(user_id=test_user_id, role="user") + + # Create UserAPIKeyAuth object + user_api_key_dict = UserAPIKeyAuth( + user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + # Mock the prisma client + mock_prisma_client = AsyncMock() + + # Mock the user table upsert operation + mock_user_response = MagicMock() + mock_user_response.model_dump.return_value = { + "user_id": test_user_id, + "user_email": None, + "teams": [test_team_id], + "user_role": "internal_user", + } + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( + return_value=mock_user_response + ) + + # Mock the budget table creation + mock_budget_response = MagicMock() + mock_budget_response.budget_id = test_new_budget_id + mock_prisma_client.db.litellm_budgettable.create = AsyncMock( + return_value=mock_budget_response + ) + + # Mock the team membership creation + mock_team_membership_response = MagicMock() + mock_team_membership_response.model_dump.return_value = { + "team_id": test_team_id, + "user_id": test_user_id, + "budget_id": test_new_budget_id, + "litellm_budget_table": None, + } + mock_prisma_client.db.litellm_teammembership.create = AsyncMock( + return_value=mock_team_membership_response + ) + + # Call the function with max_budget_in_team provided + result_user, result_team_membership = await add_new_member( + new_member=new_member, + max_budget_in_team=test_max_budget, # This should trigger budget creation + prisma_client=mock_prisma_client, + team_id=test_team_id, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=test_admin_name, + default_team_budget_id=None, # Should be ignored since max_budget_in_team is provided + ) + + # Verify that the budget was created + mock_prisma_client.db.litellm_budgettable.create.assert_called_once() + budget_call_args = mock_prisma_client.db.litellm_budgettable.create.call_args + budget_data = budget_call_args.kwargs["data"] + assert budget_data["max_budget"] == test_max_budget + assert budget_data["created_by"] == user_api_key_dict.user_id + assert budget_data["updated_by"] == user_api_key_dict.user_id + + # Verify that the team membership was created with the new budget_id + assert result_team_membership is not None + assert result_team_membership.budget_id == test_new_budget_id + + # Verify the team membership was created with the correct budget_id + team_membership_call_args = ( + mock_prisma_client.db.litellm_teammembership.create.call_args + ) + assert team_membership_call_args is not None + create_data = team_membership_call_args.kwargs["data"] + assert create_data["budget_id"] == test_new_budget_id + + +@pytest.mark.asyncio +async def test_add_new_member_with_user_email(): + """ + Test add_new_member with user_email instead of user_id and default budget. + + This test verifies that: + 1. When new_member has user_email instead of user_id + 2. And max_budget_in_team is None + 3. The default_team_budget_id is used correctly + """ + from litellm.proxy._types import LitellmUserRoles + + # Setup test data + test_user_email = "test@example.com" + test_team_id = "test_team_456" + test_default_budget_id = "default_budget_789" + test_admin_name = "test_admin" + + # Create a Member object with user_email + new_member = Member(user_email=test_user_email, role="user") + + # Create UserAPIKeyAuth object + user_api_key_dict = UserAPIKeyAuth( + user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + # Mock the prisma client + mock_prisma_client = AsyncMock() + + # Mock get_data to return empty list (no existing user) + mock_prisma_client.get_data = AsyncMock(return_value=[]) + + # Mock insert_data for new user creation + mock_user_response = MagicMock() + mock_user_response.model_dump.return_value = { + "user_id": "generated_user_id", + "user_email": test_user_email, + "teams": [test_team_id], + "user_role": "internal_user", + } + mock_prisma_client.insert_data = AsyncMock(return_value=mock_user_response) + + # Mock the team membership creation + mock_team_membership_response = MagicMock() + mock_team_membership_response.model_dump.return_value = { + "team_id": test_team_id, + "user_id": "generated_user_id", + "budget_id": test_default_budget_id, + "litellm_budget_table": None, + } + mock_prisma_client.db.litellm_teammembership.create = AsyncMock( + return_value=mock_team_membership_response + ) + + # Call the function + result_user, result_team_membership = await add_new_member( + new_member=new_member, + max_budget_in_team=None, + prisma_client=mock_prisma_client, + team_id=test_team_id, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=test_admin_name, + default_team_budget_id=test_default_budget_id, + ) + + # Verify that the user was created correctly + assert result_user is not None + assert result_user.user_email == test_user_email + + # Verify that the team membership was created with the default budget_id + assert result_team_membership is not None + assert result_team_membership.budget_id == test_default_budget_id + + # Verify that get_data was called to check for existing user + mock_prisma_client.get_data.assert_called_once_with( + key_val={"user_email": test_user_email}, + table_name="user", + query_type="find_all", + ) + + # Verify that insert_data was called to create new user + mock_prisma_client.insert_data.assert_called_once() + insert_call_args = mock_prisma_client.insert_data.call_args + insert_data = insert_call_args.kwargs["data"] + assert insert_data["user_email"] == test_user_email + assert insert_data["teams"] == [test_team_id]