tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-07 04:24:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

Revert "[Fix] Update tests for _get_and_validate_existing_key refactor"

Browse Source 
This reverts commit cf63697703.
This commit is contained in:
yuneng-jiang
2026-03-19 14:32:14 -07:00
parent cf63697703
commit e86ca7f34d
2 changed files with 107 additions and 128 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/management_endpoints/key_management_endpoints.py
-2
View File
@@ -2324,8 +2324,6 @@ async def bulk_update_keys(
error_message = error_detail.get("error", str(e))
else:
error_message = str(error_detail)
elif isinstance(e, ProxyException):
error_message = e.message
else:
error_message = str(e)
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
+107 -126
View File
@@ -5108,8 +5108,6 @@ async def test_get_and_validate_existing_key():
"""
from fastapi import HTTPException
from litellm.proxy.utils import ProxyException
# Test Case 1: Successfully retrieve existing key
mock_prisma_client = AsyncMock()
mock_key = LiteLLM_VerificationToken(
@@ -5118,37 +5116,31 @@ async def test_get_and_validate_existing_key():
models=["gpt-4"],
team_id=None,
)
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=mock_key)
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed",
return_value="hashed-test-key-123",
):
result = await _get_and_validate_existing_key(
token="test-key-123",
mock_prisma_client.get_data = AsyncMock(return_value=mock_key)
result = await _get_and_validate_existing_key(
token="test-key-123",
prisma_client=mock_prisma_client,
)
assert result == mock_key
mock_prisma_client.get_data.assert_called_once_with(
token="test-key-123",
table_name="key",
query_type="find_unique",
)
# Test Case 2: Key not found raises HTTPException
mock_prisma_client.get_data = AsyncMock(return_value=None)
with pytest.raises(HTTPException) as exc_info:
await _get_and_validate_existing_key(
token="non-existent-key",
prisma_client=mock_prisma_client,
)
assert result == mock_key
mock_prisma_client.db.litellm_verificationtoken.find_unique.assert_called_once_with(
where={"token": "hashed-test-key-123"}
)
# Test Case 2: Key not found raises ProxyException
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None)
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed",
return_value="hashed-non-existent-key",
):
with pytest.raises(ProxyException) as exc_info:
await _get_and_validate_existing_key(
token="non-existent-key",
prisma_client=mock_prisma_client,
)
assert str(exc_info.value.code) == "404"
assert "Key not found" in str(exc_info.value.message)
assert exc_info.value.status_code == 404
assert "Key not found" in str(exc_info.value.detail)
# Test Case 3: Database not connected raises HTTPException
with pytest.raises(HTTPException) as exc_info:
@@ -5197,80 +5189,75 @@ async def test_process_single_key_update():
"tags": ["production"],
}
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=existing_key)
mock_prisma_client.get_data = AsyncMock(return_value=existing_key)
mock_updated_key_obj = MagicMock()
mock_updated_key_obj.model_dump.return_value = updated_key_data
mock_prisma_client.update_data = AsyncMock(
return_value={"data": mock_updated_key_obj}
)
# Mock prepare_key_update_data
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.prepare_key_update_data"
) as mock_prepare:
mock_prepare.return_value = {"max_budget": 100.0, "tags": ["production"]}
# Mock TeamMemberPermissionChecks
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint"
) as mock_permission_check:
mock_permission_check.return_value = None
# Mock _delete_cache_key_object
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object"
) as mock_delete_cache:
mock_delete_cache.return_value = None
# Mock hash_token (imported from litellm.proxy._types)
with patch(
"litellm.proxy._types.hash_token"
) as mock_hash:
mock_hash.return_value = "hashed-test-key-123"
# Mock _hash_token_if_needed
# Mock KeyManagementEventHooks
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed",
return_value="hashed-test-key-123",
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook"
):
# Mock KeyManagementEventHooks
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook"
):
# Create update request
key_update_item = BulkUpdateKeyRequestItem(
key="test-key-123",
max_budget=100.0,
tags=["production"],
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-admin",
user_id="admin-user",
)
# Call the function
result = await _process_single_key_update(
key_update_item=key_update_item,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
prisma_client=mock_prisma_client,
user_api_key_cache=mock_user_api_key_cache,
proxy_logging_obj=mock_proxy_logging_obj,
llm_router=mock_llm_router,
)
# Verify results
assert result is not None
assert "token" not in result # Token should be removed
assert result.get("max_budget") == 100.0
assert result.get("tags") == ["production"]
# Verify mocks were called
mock_prisma_client.db.litellm_verificationtoken.find_unique.assert_called_once()
mock_prisma_client.update_data.assert_called_once()
mock_delete_cache.assert_called_once()
# Create update request
key_update_item = BulkUpdateKeyRequestItem(
key="test-key-123",
max_budget=100.0,
tags=["production"],
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-admin",
user_id="admin-user",
)
# Call the function
result = await _process_single_key_update(
key_update_item=key_update_item,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
prisma_client=mock_prisma_client,
user_api_key_cache=mock_user_api_key_cache,
proxy_logging_obj=mock_proxy_logging_obj,
llm_router=mock_llm_router,
)
# Verify results
assert result is not None
assert "token" not in result # Token should be removed
assert result.get("max_budget") == 100.0
assert result.get("tags") == ["production"]
# Verify mocks were called
mock_prisma_client.get_data.assert_called_once()
mock_prisma_client.update_data.assert_called_once()
mock_delete_cache.assert_called_once()
@pytest.mark.asyncio
@@ -5460,17 +5447,15 @@ async def test_bulk_update_keys_partial_failures(monkeypatch):
}
# First key exists, second key doesn't exist
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
mock_prisma_client.get_data = AsyncMock(
side_effect=[existing_key_1, None] # Second key not found
)
# get_data is used by the error handler to fetch key info for failed updates
mock_prisma_client.get_data = AsyncMock(return_value=None)
mock_updated_key_1_obj = MagicMock()
mock_updated_key_1_obj.model_dump.return_value = updated_key_1_data
mock_prisma_client.update_data = AsyncMock(
return_value={"data": mock_updated_key_1_obj}
)
# Patch dependencies
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client", mock_prisma_client
@@ -5482,13 +5467,13 @@ async def test_bulk_update_keys_partial_failures(monkeypatch):
"litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj
)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_llm_router)
# Mock helper functions
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.prepare_key_update_data"
) as mock_prepare:
mock_prepare.return_value = {"max_budget": 100.0, "tags": ["production"]}
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint"
):
@@ -5499,50 +5484,46 @@ async def test_bulk_update_keys_partial_failures(monkeypatch):
"litellm.proxy._types.hash_token"
) as mock_hash:
mock_hash.return_value = "hashed-key-1"
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed",
side_effect=lambda token: f"hashed-{token}",
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook"
):
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook"
):
# Create request with one valid and one invalid key
request_data = BulkUpdateKeyRequest(
keys=[
BulkUpdateKeyRequestItem(
key="test-key-1",
max_budget=100.0,
tags=["production"],
),
BulkUpdateKeyRequestItem(
key="non-existent-key",
max_budget=200.0,
tags=["staging"],
),
]
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-admin",
user_id="admin-user",
)
# Call endpoint
response = await bulk_update_keys(
data=request_data,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
)
# Verify response
assert response.total_requested == 2
assert len(response.successful_updates) == 1
assert len(response.failed_updates) == 1
assert response.successful_updates[0].key == "test-key-1"
assert response.failed_updates[0].key == "non-existent-key"
assert "Key not found" in response.failed_updates[0].failed_reason
# Create request with one valid and one invalid key
request_data = BulkUpdateKeyRequest(
keys=[
BulkUpdateKeyRequestItem(
key="test-key-1",
max_budget=100.0,
tags=["production"],
),
BulkUpdateKeyRequestItem(
key="non-existent-key",
max_budget=200.0,
tags=["staging"],
),
]
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-admin",
user_id="admin-user",
)
# Call endpoint
response = await bulk_update_keys(
data=request_data,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
)
# Verify response
assert response.total_requested == 2
assert len(response.successful_updates) == 1
assert len(response.failed_updates) == 1
assert response.successful_updates[0].key == "test-key-1"
assert response.failed_updates[0].key == "non-existent-key"
assert "Key not found" in response.failed_updates[0].failed_reason
@pytest.mark.parametrize(