fix test_async_log_proxy_authentication_errors_get_request

This commit is contained in:
Ishaan Jaff 2024-11-27 11:58:09 -08:00
parent c09ca785a4
commit cc83463707

View file

@ -2195,3 +2195,126 @@ async def test_async_log_proxy_authentication_errors():
assert (
mock_logger.user_api_key_dict_logged.token is not None
) # token should be hashed
@pytest.mark.asyncio
async def test_async_log_proxy_authentication_errors_get_request():
"""
Test if async_log_proxy_authentication_errors correctly handles GET requests
that don't have a JSON body
"""
import json
from fastapi import Request
from litellm.proxy.utils import ProxyLogging
from litellm.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
class MockCustomLogger(CustomLogger):
def __init__(self):
self.called = False
self.exception_logged = None
self.request_data_logged = None
self.user_api_key_dict_logged = None
async def async_post_call_failure_hook(
self,
request_data: dict,
original_exception: Exception,
user_api_key_dict: UserAPIKeyAuth,
):
self.called = True
self.exception_logged = original_exception
self.request_data_logged = request_data
self.user_api_key_dict_logged = user_api_key_dict
# Create a mock GET request
request = Request(scope={"type": "http", "method": "GET"})
# Mock the json() method to raise JSONDecodeError
async def mock_json():
raise json.JSONDecodeError("Expecting value", "", 0)
request.json = mock_json
# Create a test exception
test_exception = Exception("Invalid API Key")
# Initialize ProxyLogging
mock_logger = MockCustomLogger()
litellm.callbacks = [mock_logger]
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
# Call the method
await proxy_logging_obj.async_log_proxy_authentication_errors(
original_exception=test_exception,
request=request,
parent_otel_span=None,
api_key="test-key",
)
# Verify the mock logger was called with correct parameters
assert mock_logger.called == True
assert mock_logger.exception_logged == test_exception
assert mock_logger.user_api_key_dict_logged is not None
assert mock_logger.user_api_key_dict_logged.token is not None
@pytest.mark.asyncio
async def test_async_log_proxy_authentication_errors_no_api_key():
"""
Test if async_log_proxy_authentication_errors correctly handles requests
with no API key provided
"""
from fastapi import Request
from litellm.proxy.utils import ProxyLogging
from litellm.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
class MockCustomLogger(CustomLogger):
def __init__(self):
self.called = False
self.exception_logged = None
self.request_data_logged = None
self.user_api_key_dict_logged = None
async def async_post_call_failure_hook(
self,
request_data: dict,
original_exception: Exception,
user_api_key_dict: UserAPIKeyAuth,
):
self.called = True
self.exception_logged = original_exception
self.request_data_logged = request_data
self.user_api_key_dict_logged = user_api_key_dict
# Create test data
test_data = {"model": "gpt-4", "messages": [{"role": "user", "content": "Hello"}]}
# Create a mock request
request = Request(scope={"type": "http", "method": "POST"})
request._json = AsyncMock(return_value=test_data)
# Create a test exception
test_exception = Exception("No API Key Provided")
# Initialize ProxyLogging
mock_logger = MockCustomLogger()
litellm.callbacks = [mock_logger]
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
# Call the method with api_key=None
await proxy_logging_obj.async_log_proxy_authentication_errors(
original_exception=test_exception,
request=request,
parent_otel_span=None,
api_key=None,
)
# Verify the mock logger was called with correct parameters
assert mock_logger.called == True
assert mock_logger.exception_logged == test_exception
assert mock_logger.user_api_key_dict_logged is not None
assert (
mock_logger.user_api_key_dict_logged.token == ""
) # Empty token for no API key