mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
test: add unit testing for both flows on key unblock
This commit is contained in:
Krrish Dholakia
parent
0f6898ad0a
commit
0c8b311155
1 changed files with 202 additions and 43 deletions
|
|
@ -183,7 +183,9 @@ async def test_budget_reset_and_expires_at_first_of_month(monkeypatch):
|
|||
assert (
|
||||
response_date.month == expected_month
|
||||
), f"Expected month {expected_month}, got {response_date.month} for {key}"
|
||||
assert response_date.day == 1, f"Expected day 1, got {response_date.day} for {key}"
|
||||
assert (
|
||||
response_date.day == 1
|
||||
), f"Expected day 1, got {response_date.day} for {key}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -507,7 +509,6 @@ def test_get_new_token_with_invalid_key():
|
|||
assert "New key must start with 'sk-'" in str(exc_info.value.detail)
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_service_account_requires_team_id():
|
||||
with pytest.raises(HTTPException):
|
||||
|
|
@ -529,11 +530,12 @@ async def test_generate_service_account_works_with_team_id():
|
|||
from unittest.mock import patch
|
||||
|
||||
# Mock the database and router dependencies from proxy_server
|
||||
with patch('litellm.proxy.proxy_server.prisma_client') as mock_prisma, \
|
||||
patch('litellm.proxy.proxy_server.llm_router') as mock_router, \
|
||||
patch('litellm.proxy.proxy_server.premium_user', False), \
|
||||
patch('litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn') as mock_generate_key:
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch(
|
||||
"litellm.proxy.proxy_server.llm_router"
|
||||
) as mock_router, patch("litellm.proxy.proxy_server.premium_user", False), patch(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn"
|
||||
) as mock_generate_key:
|
||||
|
||||
# Configure mocks
|
||||
mock_prisma.return_value = AsyncMock()
|
||||
mock_router.return_value = None
|
||||
|
|
@ -542,9 +544,9 @@ async def test_generate_service_account_works_with_team_id():
|
|||
"key": "sk-test-key",
|
||||
"expires": None,
|
||||
"user_id": "test-user",
|
||||
"team_id": "IJ"
|
||||
"team_id": "IJ",
|
||||
}
|
||||
|
||||
|
||||
# This should not raise an exception since team_id is provided
|
||||
await _common_key_generation_helper(
|
||||
data=GenerateKeyRequest(
|
||||
|
|
@ -559,7 +561,6 @@ async def test_generate_service_account_works_with_team_id():
|
|||
)
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_service_account_requires_team_id():
|
||||
data = UpdateKeyRequest(key="sk-1", metadata={"service_account_id": "sa"})
|
||||
|
|
@ -571,7 +572,9 @@ async def test_update_service_account_requires_team_id():
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_service_account_works_with_team_id():
|
||||
data = UpdateKeyRequest(key="sk-1", metadata={"service_account_id": "sa"}, team_id="IJ")
|
||||
data = UpdateKeyRequest(
|
||||
key="sk-1", metadata={"service_account_id": "sa"}, team_id="IJ"
|
||||
)
|
||||
existing_key = LiteLLM_VerificationToken(token="hashed")
|
||||
|
||||
await prepare_key_update_data(data=data, existing_key_row=existing_key)
|
||||
|
|
@ -580,22 +583,22 @@ async def test_update_service_account_works_with_team_id():
|
|||
@pytest.mark.asyncio
|
||||
async def test_validate_team_id_used_in_service_account_request_requires_team_id():
|
||||
"""
|
||||
Test that validate_team_id_used_in_service_account_request raises HTTPException
|
||||
Test that validate_team_id_used_in_service_account_request raises HTTPException
|
||||
when team_id is None for service account key generation.
|
||||
"""
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
validate_team_id_used_in_service_account_request,
|
||||
)
|
||||
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
|
||||
|
||||
# Test that HTTPException is raised when team_id is None
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await validate_team_id_used_in_service_account_request(
|
||||
team_id=None,
|
||||
prisma_client=mock_prisma_client,
|
||||
)
|
||||
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "team_id is required for service account keys" in str(exc_info.value.detail)
|
||||
|
||||
|
|
@ -603,7 +606,7 @@ async def test_validate_team_id_used_in_service_account_request_requires_team_id
|
|||
@pytest.mark.asyncio
|
||||
async def test_validate_team_id_used_in_service_account_request_requires_prisma_client():
|
||||
"""
|
||||
Test that validate_team_id_used_in_service_account_request raises HTTPException
|
||||
Test that validate_team_id_used_in_service_account_request raises HTTPException
|
||||
when prisma_client is None for service account key generation.
|
||||
"""
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
|
|
@ -616,78 +619,76 @@ async def test_validate_team_id_used_in_service_account_request_requires_prisma_
|
|||
team_id="test-team-id",
|
||||
prisma_client=None,
|
||||
)
|
||||
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "prisma_client is required for service account keys" in str(exc_info.value.detail)
|
||||
assert "prisma_client is required for service account keys" in str(
|
||||
exc_info.value.detail
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_validate_team_id_used_in_service_account_request_checks_team_exists():
|
||||
"""
|
||||
Test that validate_team_id_used_in_service_account_request validates that
|
||||
Test that validate_team_id_used_in_service_account_request validates that
|
||||
the team_id exists in the database for service account key generation.
|
||||
"""
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
validate_team_id_used_in_service_account_request,
|
||||
)
|
||||
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
|
||||
|
||||
# Mock the database query to return None (team doesn't exist)
|
||||
mock_find_unique = AsyncMock(return_value=None)
|
||||
mock_prisma_client.db.litellm_teamtable.find_unique = mock_find_unique
|
||||
|
||||
|
||||
# Test that HTTPException is raised when team doesn't exist in DB
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await validate_team_id_used_in_service_account_request(
|
||||
team_id="non-existent-team-id",
|
||||
prisma_client=mock_prisma_client,
|
||||
)
|
||||
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "team_id does not exist in the database" in str(exc_info.value.detail)
|
||||
|
||||
|
||||
# Verify the database was queried with the correct parameters
|
||||
mock_find_unique.assert_called_once_with(
|
||||
where={"team_id": "non-existent-team-id"}
|
||||
)
|
||||
mock_find_unique.assert_called_once_with(where={"team_id": "non-existent-team-id"})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_validate_team_id_used_in_service_account_request_success():
|
||||
"""
|
||||
Test that validate_team_id_used_in_service_account_request returns True
|
||||
Test that validate_team_id_used_in_service_account_request returns True
|
||||
when team_id exists in the database for service account key generation.
|
||||
"""
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
validate_team_id_used_in_service_account_request,
|
||||
)
|
||||
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
|
||||
|
||||
# Mock the database query to return a team object (team exists)
|
||||
mock_team = {"team_id": "existing-team-id", "team_name": "Test Team"}
|
||||
mock_find_unique = AsyncMock(return_value=mock_team)
|
||||
mock_prisma_client.db.litellm_teamtable.find_unique = mock_find_unique
|
||||
|
||||
|
||||
# Test that function returns True when team exists
|
||||
result = await validate_team_id_used_in_service_account_request(
|
||||
team_id="existing-team-id",
|
||||
prisma_client=mock_prisma_client,
|
||||
)
|
||||
|
||||
|
||||
assert result is True
|
||||
|
||||
|
||||
# Verify the database was queried with the correct parameters
|
||||
mock_find_unique.assert_called_once_with(
|
||||
where={"team_id": "existing-team-id"}
|
||||
)
|
||||
mock_find_unique.assert_called_once_with(where={"team_id": "existing-team-id"})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_service_account_key_endpoint_validation():
|
||||
"""
|
||||
Test that the /key/service-account/generate endpoint properly validates
|
||||
Test that the /key/service-account/generate endpoint properly validates
|
||||
team_id requirement and team existence in database.
|
||||
"""
|
||||
from unittest.mock import patch
|
||||
|
|
@ -705,16 +706,16 @@ async def test_generate_service_account_key_endpoint_validation():
|
|||
),
|
||||
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:
|
||||
|
||||
# Test case 2: Team doesn't exist in database
|
||||
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
|
||||
# Mock team not found
|
||||
mock_find_unique = AsyncMock(return_value=None)
|
||||
mock_prisma.db.litellm_teamtable.find_unique = mock_find_unique
|
||||
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await generate_service_account_key_fn(
|
||||
data=GenerateKeyRequest(team_id="non-existent-team"),
|
||||
|
|
@ -723,7 +724,165 @@ async def test_generate_service_account_key_endpoint_validation():
|
|||
),
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "team_id does not exist in the database" in str(exc_info.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unblock_key_supports_both_sk_and_hashed_tokens(monkeypatch):
|
||||
"""
|
||||
Test that the unblock_key endpoint correctly handles both sk- prefixed tokens
|
||||
and hashed tokens by properly converting sk- tokens to hashed format before
|
||||
database operations.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.proxy._types import BlockKeyRequest
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import unblock_key
|
||||
|
||||
# Mock dependencies
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_user_api_key_cache = MagicMock()
|
||||
mock_proxy_logging_obj = MagicMock()
|
||||
|
||||
# Use a proper 64-character hex hash for testing
|
||||
test_hashed_token = (
|
||||
"a1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd"
|
||||
)
|
||||
|
||||
# Mock the key record that will be returned from database
|
||||
mock_key_record = MagicMock()
|
||||
mock_key_record.token = test_hashed_token
|
||||
mock_key_record.blocked = False
|
||||
mock_key_record.model_dump_json.return_value = (
|
||||
f'{{"token": "{test_hashed_token}", "blocked": false}}'
|
||||
)
|
||||
|
||||
# Mock database operations
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
|
||||
return_value=mock_key_record
|
||||
)
|
||||
mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock(
|
||||
return_value=mock_key_record
|
||||
)
|
||||
|
||||
# Mock get_key_object and _cache_key_object functions
|
||||
mock_key_object = MagicMock()
|
||||
mock_key_object.blocked = True # Initially blocked
|
||||
|
||||
# Mock hash_token function
|
||||
def mock_hash_token(token):
|
||||
if token == "sk-test123456789":
|
||||
return test_hashed_token
|
||||
return token
|
||||
|
||||
# Apply monkeypatch
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.hash_token", mock_hash_token)
|
||||
monkeypatch.setattr(
|
||||
"litellm.store_audit_logs", False
|
||||
) # Disable audit logs for simpler test
|
||||
|
||||
# Mock get_key_object and _cache_key_object
|
||||
async def mock_get_key_object(**kwargs):
|
||||
return mock_key_object
|
||||
|
||||
async def mock_cache_key_object(**kwargs):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.get_key_object",
|
||||
mock_get_key_object,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints._cache_key_object",
|
||||
mock_cache_key_object,
|
||||
)
|
||||
|
||||
# Create mock request and user auth
|
||||
mock_request = MagicMock()
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin_user"
|
||||
)
|
||||
|
||||
# Test Case 1: Using sk- prefixed token
|
||||
sk_token_request = BlockKeyRequest(key="sk-test123456789")
|
||||
|
||||
result = await unblock_key(
|
||||
data=sk_token_request,
|
||||
http_request=mock_request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
# Verify that the database update was called with hashed token
|
||||
mock_prisma_client.db.litellm_verificationtoken.update.assert_called_with(
|
||||
where={"token": test_hashed_token}, data={"blocked": False}
|
||||
)
|
||||
|
||||
assert result == mock_key_record
|
||||
assert mock_key_object.blocked == False # Should be updated to unblocked
|
||||
|
||||
# Reset mocks for second test
|
||||
mock_prisma_client.db.litellm_verificationtoken.update.reset_mock()
|
||||
mock_key_object.blocked = True # Reset to blocked state
|
||||
|
||||
# Test Case 2: Using already hashed token
|
||||
hashed_token_request = BlockKeyRequest(key=test_hashed_token)
|
||||
|
||||
result = await unblock_key(
|
||||
data=hashed_token_request,
|
||||
http_request=mock_request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
# Verify that the database update was called with the same hashed token
|
||||
mock_prisma_client.db.litellm_verificationtoken.update.assert_called_with(
|
||||
where={"token": test_hashed_token}, data={"blocked": False}
|
||||
)
|
||||
|
||||
assert result == mock_key_record
|
||||
assert mock_key_object.blocked == False # Should be updated to unblocked
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unblock_key_invalid_key_format(monkeypatch):
|
||||
"""
|
||||
Test that unblock_key properly validates key format and raises appropriate errors
|
||||
for invalid keys.
|
||||
"""
|
||||
from litellm.proxy._types import BlockKeyRequest
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import unblock_key
|
||||
from litellm.proxy.utils import ProxyException
|
||||
|
||||
# Mock prisma_client to avoid DB connection error
|
||||
mock_prisma_client = AsyncMock()
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
|
||||
# Mock request and user auth
|
||||
mock_request = MagicMock()
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin_user"
|
||||
)
|
||||
|
||||
# Test with invalid key format
|
||||
invalid_key_request = BlockKeyRequest(key="invalid-key-format")
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await unblock_key(
|
||||
data=invalid_key_request,
|
||||
http_request=mock_request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
assert exc_info.value.code == "400"
|
||||
assert "Invalid key format" in str(exc_info.value.message)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue