From 814e49df0691b0cda85a753eddf7693582c0dd16 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 4 Oct 2025 14:39:50 -0700 Subject: [PATCH] test: update tests --- .../test_key_management_endpoints.py | 35 +++++++++++++------ 1 file changed, 25 insertions(+), 10 deletions(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 0b659acd2bb..7620e27e4c6 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -850,17 +850,24 @@ async def test_generate_service_account_key_endpoint_validation(): ) # Test case 1: Missing team_id - with pytest.raises(HTTPException) as exc_info: - await generate_service_account_key_fn( - data=GenerateKeyRequest(team_id=None), - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1" - ), - litellm_changed_by=None, - ) + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + # Mock prisma_client to be not None so we can reach team_id validation + mock_prisma_instance = AsyncMock() + mock_prisma.return_value = mock_prisma_instance - assert exc_info.value.status_code == 400 - assert "team_id is required for service account keys" in str(exc_info.value.detail) + with pytest.raises(HTTPException) as exc_info: + await generate_service_account_key_fn( + data=GenerateKeyRequest(team_id=None), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1" + ), + litellm_changed_by=None, + ) + + assert exc_info.value.status_code == 400 + assert "team_id is required for service account keys" in str( + exc_info.value.detail + ) # Test case 2: Team doesn't exist in database with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: @@ -1257,6 +1264,8 @@ async def test_check_team_key_limits_no_existing_keys(): data = GenerateKeyRequest( tpm_limit=5000, rpm_limit=500, + tpm_limit_type="guaranteed_throughput", + rpm_limit_type="guaranteed_throughput", ) # Should not raise any exception @@ -1333,10 +1342,12 @@ async def test_check_team_key_limits_tpm_overallocation(): existing_key1 = MagicMock() existing_key1.tpm_limit = 6000 existing_key1.rpm_limit = 100 + existing_key1.metadata = {} existing_key2 = MagicMock() existing_key2.tpm_limit = 3000 existing_key2.rpm_limit = 200 + existing_key2.metadata = {} mock_prisma_client = AsyncMock() mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( @@ -1360,6 +1371,7 @@ async def test_check_team_key_limits_tpm_overallocation(): data = GenerateKeyRequest( tpm_limit=2000, rpm_limit=100, + tpm_limit_type="guaranteed_throughput", ) # Should raise HTTPException for TPM overallocation @@ -1387,10 +1399,12 @@ async def test_check_team_key_limits_rpm_overallocation(): existing_key1 = MagicMock() existing_key1.tpm_limit = 1000 existing_key1.rpm_limit = 600 + existing_key1.metadata = {} existing_key2 = MagicMock() existing_key2.tpm_limit = 2000 existing_key2.rpm_limit = 300 + existing_key2.metadata = {} mock_prisma_client = AsyncMock() mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( @@ -1414,6 +1428,7 @@ async def test_check_team_key_limits_rpm_overallocation(): data = GenerateKeyRequest( tpm_limit=1000, rpm_limit=200, + rpm_limit_type="guaranteed_throughput", ) # Should raise HTTPException for RPM overallocation