diff --git a/tests/test_litellm/integrations/test_prometheus_invalid_key_filtering.py b/tests/test_litellm/integrations/test_prometheus_invalid_key_filtering.py new file mode 100644 index 00000000000..ff433480d5e --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_invalid_key_filtering.py @@ -0,0 +1,161 @@ +""" +Unit tests for Prometheus invalid API key request filtering. + +Tests functionality that prevents invalid API key requests (401 status codes) +from being recorded in Prometheus metrics. +""" + +import os +import sys +from unittest.mock import Mock, patch + +import pytest +from prometheus_client import REGISTRY + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.integrations.prometheus import PrometheusLogger +from litellm.proxy._types import UserAPIKeyAuth + + +@pytest.fixture(scope="function") +def prometheus_logger(): + """Create a PrometheusLogger instance for testing.""" + collectors = list(REGISTRY._collector_to_names.keys()) + for collector in collectors: + REGISTRY.unregister(collector) + return PrometheusLogger() + + +class ExceptionWithCode: + """Exception-like object with 'code' attribute (ProxyException pattern).""" + def __init__(self, code): + self.code = code + + +class ExceptionWithStatusCode: + """Exception-like object with 'status_code' attribute.""" + def __init__(self, status_code): + self.status_code = status_code + + +class TestExtractStatusCode: + """Test status code extraction from various sources.""" + + @pytest.mark.parametrize("exception_class,code_value,expected", [ + (ExceptionWithCode, "401", 401), + (ExceptionWithStatusCode, 401, 401), + ]) + def test_extract_from_exception(self, prometheus_logger, exception_class, code_value, expected): + exception = exception_class(code_value) + assert prometheus_logger._extract_status_code(exception=exception) == expected + + def test_extract_from_kwargs(self, prometheus_logger): + exception = ExceptionWithCode("401") + assert prometheus_logger._extract_status_code(kwargs={"exception": exception}) == 401 + + def test_extract_from_enum_values(self, prometheus_logger): + enum_values = Mock(status_code="401") + assert prometheus_logger._extract_status_code(enum_values=enum_values) == 401 + + +class TestInvalidAPIKeyDetection: + """Test invalid API key request detection logic.""" + + @pytest.mark.parametrize("status_code,expected", [ + (401, True), + (200, False), + (500, False), + (None, False), + ]) + def test_status_code_detection(self, prometheus_logger, status_code, expected): + assert prometheus_logger._is_invalid_api_key_request(status_code=status_code) == expected + + def test_auth_error_message_detection(self, prometheus_logger): + exception = AssertionError("LiteLLM Virtual Key expected. Received=invalid-key-12345, expected to start with 'sk-'.") + assert prometheus_logger._is_invalid_api_key_request(status_code=None, exception=exception) is True + + def test_non_auth_exception_not_detected(self, prometheus_logger): + exception = ValueError("Some other error") + assert prometheus_logger._is_invalid_api_key_request(status_code=None, exception=exception) is False + + +class TestSkipMetricsValidation: + """Test high-level validation method that orchestrates detection and extraction.""" + + def test_skip_for_401_exception(self, prometheus_logger): + """Test full flow: extraction -> detection -> skip decision.""" + exception = ExceptionWithCode("401") + assert prometheus_logger._should_skip_metrics_for_invalid_key(exception=exception) is True + + def test_skip_for_auth_error_message(self, prometheus_logger): + """Test full flow: exception message -> detection -> skip decision.""" + exception = AssertionError("expected to start with 'sk-'") + assert prometheus_logger._should_skip_metrics_for_invalid_key(exception=exception) is True + + def test_no_skip_for_valid_request(self, prometheus_logger): + assert prometheus_logger._should_skip_metrics_for_invalid_key() is False + + +class TestAsyncHooks: + """Test async hook methods skip metrics for invalid API keys.""" + + @pytest.fixture + def mock_user_api_key(self): + """Create a mock UserAPIKeyAuth object.""" + user_key = Mock(spec=UserAPIKeyAuth) + user_key.api_key = "test-key" + user_key.end_user_id = None + user_key.user_id = None + user_key.user_email = None + user_key.key_alias = None + user_key.team_id = None + user_key.team_alias = None + user_key.request_route = "/test" + return user_key + + @pytest.mark.asyncio + async def test_post_call_failure_hook_skips_401(self, prometheus_logger, mock_user_api_key): + exception = ExceptionWithCode("401") + exception.__class__.__name__ = "ProxyException" + + with patch.object(prometheus_logger, 'litellm_proxy_failed_requests_metric') as mock_failed, \ + patch.object(prometheus_logger, 'litellm_proxy_total_requests_metric') as mock_total: + + await prometheus_logger.async_post_call_failure_hook( + request_data={"model": "test-model"}, + original_exception=exception, + user_api_key_dict=mock_user_api_key + ) + + mock_failed.labels.assert_not_called() + mock_total.labels.assert_not_called() + + @pytest.mark.asyncio + async def test_log_failure_event_skips_401(self, prometheus_logger): + exception = ExceptionWithCode("401") + kwargs = { + "model": "test-model", + "standard_logging_object": { + "metadata": { + "user_api_key_hash": "test-key", + "user_api_key_user_id": "test-user", + }, + "model_group": "test-model", + }, + "exception": exception, + "litellm_params": {}, + } + + with patch.object(prometheus_logger, 'litellm_llm_api_failed_requests_metric') as mock_failed, \ + patch.object(prometheus_logger, 'set_llm_deployment_failure_metrics') as mock_deployment: + + await prometheus_logger.async_log_failure_event( + kwargs=kwargs, + response_obj=None, + start_time=None, + end_time=None + ) + + mock_failed.labels.assert_not_called() + mock_deployment.assert_not_called()