scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

test: update tests

Browse source
This commit is contained in:
Krrish Dholakia 2025-10-04 14:39:50 -07:00
parent f500a866dd
commit 814e49df06
1 changed files with 25 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

35
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -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