mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-11 11:04:32 +00:00
Mateo Wang
bfbb5d2375
* test(logging): align DB metrics event_metadata assertions with safe redaction PR #28909 hardened log_db_metrics to emit a minimal, non-sensitive event_metadata (only table_name when present, otherwise None) instead of dumping function_name, function_kwargs, and function_args onto the span. The test in test_log_db_redis_services was not updated and still asserted "function_name" in event_metadata, which raised TypeError (argument of type 'NoneType' is not iterable) and turned the logging_testing CI job red on litellm_internal_staging. Update test_log_db_metrics_success to assert event_metadata is None when no table_name is passed, and add test_log_db_metrics_event_metadata_is_safe as a regression guard verifying that only the table name surfaces and that sensitive kwargs (tokens, prisma client) are never dumped. * test(bedrock): self-heal opus-4-7 grid cells when unentitled on CI The bedrock-claude-opus-4-7 converse cells are unentitled on the Bedrock CI account, so they were marked xfail. xfail keeps reporting them as expected failures even after access is granted, so the wire translation never gets verified again. Now the cell makes the call and skips only when Bedrock replies "is not available for this account"; the moment the model is entitled the same cells run their full assertions with no edit. A focused unit test pins the tolerance predicate so any other failure still surfaces loudly and the available path still runs the assertions.
237 lines
8.2 KiB
Python
237 lines
8.2 KiB
Python
import io
|
|
import os
|
|
import sys
|
|
|
|
sys.path.insert(0, os.path.abspath("../.."))
|
|
|
|
import asyncio
|
|
import gzip
|
|
import json
|
|
import logging
|
|
import time
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
import pytest
|
|
|
|
import litellm
|
|
from litellm import completion
|
|
from litellm._logging import verbose_logger
|
|
from litellm.proxy.utils import log_db_metrics, ServiceTypes
|
|
from datetime import datetime
|
|
import httpx
|
|
from prisma.errors import ClientNotConnectedError
|
|
|
|
|
|
# Test async function to decorate
|
|
@log_db_metrics
|
|
async def sample_db_function(*args, **kwargs):
|
|
return "success"
|
|
|
|
|
|
@log_db_metrics
|
|
async def sample_proxy_function(*args, **kwargs):
|
|
return "success"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_log_db_metrics_success():
|
|
# Mock the proxy_logging_obj
|
|
with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging:
|
|
# Setup mock
|
|
mock_proxy_logging.service_logging_obj.async_service_success_hook = AsyncMock()
|
|
|
|
# Call the decorated function
|
|
result = await sample_db_function(parent_otel_span="test_span")
|
|
|
|
# Assertions
|
|
assert result == "success"
|
|
mock_proxy_logging.service_logging_obj.async_service_success_hook.assert_called_once()
|
|
call_args = (
|
|
mock_proxy_logging.service_logging_obj.async_service_success_hook.call_args[
|
|
1
|
|
]
|
|
)
|
|
|
|
assert call_args["service"] == ServiceTypes.DB
|
|
assert call_args["call_type"] == "sample_db_function"
|
|
assert call_args["parent_otel_span"] == "test_span"
|
|
assert isinstance(call_args["duration"], float)
|
|
assert isinstance(call_args["start_time"], datetime)
|
|
assert isinstance(call_args["end_time"], datetime)
|
|
assert call_args["event_metadata"] is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_log_db_metrics_event_metadata_is_safe():
|
|
"""event_metadata must surface only the table name, never the raw
|
|
kwargs/args which carry live clients (Prisma, OTel spans) and secrets.
|
|
|
|
Regression guard for #28909: a previous version dumped function_kwargs and
|
|
function_args onto the span.
|
|
"""
|
|
with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging:
|
|
mock_proxy_logging.service_logging_obj.async_service_success_hook = AsyncMock()
|
|
|
|
@log_db_metrics
|
|
async def db_call(**kwargs):
|
|
return "success"
|
|
|
|
await db_call(
|
|
parent_otel_span="test_span",
|
|
table_name="LiteLLM_SpendLogs",
|
|
token="sk-secret-should-not-leak",
|
|
prisma_client=object(),
|
|
)
|
|
await asyncio.sleep(0)
|
|
|
|
call_args = (
|
|
mock_proxy_logging.service_logging_obj.async_service_success_hook.call_args[
|
|
1
|
|
]
|
|
)
|
|
assert call_args["event_metadata"] == {"table_name": "LiteLLM_SpendLogs"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_log_db_metrics_duration():
|
|
# Mock the proxy_logging_obj
|
|
with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging:
|
|
# Setup mock
|
|
mock_proxy_logging.service_logging_obj.async_service_success_hook = AsyncMock()
|
|
|
|
# Add a delay to the function to test duration
|
|
@log_db_metrics
|
|
async def delayed_function(**kwargs):
|
|
await asyncio.sleep(1) # 1 second delay
|
|
return "success"
|
|
|
|
# Call the decorated function
|
|
start = time.time()
|
|
result = await delayed_function(parent_otel_span="test_span")
|
|
end = time.time()
|
|
|
|
# Get the actual duration
|
|
actual_duration = end - start
|
|
|
|
# Get the logged duration from the mock call
|
|
call_args = (
|
|
mock_proxy_logging.service_logging_obj.async_service_success_hook.call_args[
|
|
1
|
|
]
|
|
)
|
|
logged_duration = call_args["duration"]
|
|
|
|
# Assert the logged duration is approximately equal to actual duration (within 0.1 seconds)
|
|
assert abs(logged_duration - actual_duration) < 0.1
|
|
assert result == "success"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_log_db_metrics_failure():
|
|
"""
|
|
should log a failure if a prisma error is raised
|
|
"""
|
|
# Mock the proxy_logging_obj
|
|
from prisma.errors import ClientNotConnectedError
|
|
|
|
with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging:
|
|
# Setup mock
|
|
mock_proxy_logging.service_logging_obj.async_service_failure_hook = AsyncMock()
|
|
|
|
# Create a failing function
|
|
@log_db_metrics
|
|
async def failing_function(**kwargs):
|
|
raise ClientNotConnectedError()
|
|
|
|
# Call the decorated function and expect it to raise
|
|
with pytest.raises(ClientNotConnectedError) as exc_info:
|
|
await failing_function(parent_otel_span="test_span")
|
|
|
|
# Assertions
|
|
assert "Client is not connected to the query engine" in str(exc_info.value)
|
|
mock_proxy_logging.service_logging_obj.async_service_failure_hook.assert_called_once()
|
|
call_args = (
|
|
mock_proxy_logging.service_logging_obj.async_service_failure_hook.call_args[
|
|
1
|
|
]
|
|
)
|
|
|
|
assert call_args["service"] == ServiceTypes.DB
|
|
assert call_args["call_type"] == "failing_function"
|
|
assert call_args["parent_otel_span"] == "test_span"
|
|
assert isinstance(call_args["duration"], float)
|
|
assert isinstance(call_args["start_time"], datetime)
|
|
assert isinstance(call_args["end_time"], datetime)
|
|
assert isinstance(call_args["error"], ClientNotConnectedError)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"exception,should_log",
|
|
[
|
|
(ValueError("Generic error"), False),
|
|
(KeyError("Missing key"), False),
|
|
(TypeError("Type error"), False),
|
|
(httpx.ConnectError("Failed to connect"), True),
|
|
(httpx.TimeoutException("Request timed out"), True),
|
|
(ClientNotConnectedError(), True), # Prisma error
|
|
],
|
|
)
|
|
async def test_log_db_metrics_failure_error_types(exception, should_log):
|
|
"""
|
|
Why Test?
|
|
Users were seeing that non-DB errors were being logged as DB Service Failures
|
|
Example a failure to read a value from cache was being logged as a DB Service Failure
|
|
|
|
|
|
Parameterized test to verify:
|
|
- DB-related errors (Prisma, httpx) are logged as service failures
|
|
- Non-DB errors (ValueError, KeyError, etc.) are not logged
|
|
"""
|
|
with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging:
|
|
mock_proxy_logging.service_logging_obj.async_service_failure_hook = AsyncMock()
|
|
|
|
@log_db_metrics
|
|
async def failing_function(**kwargs):
|
|
raise exception
|
|
|
|
# Call the function and expect it to raise the exception
|
|
with pytest.raises(type(exception)):
|
|
await failing_function(parent_otel_span="test_span")
|
|
|
|
if should_log:
|
|
# Assert failure was logged for DB-related errors
|
|
mock_proxy_logging.service_logging_obj.async_service_failure_hook.assert_called_once()
|
|
call_args = mock_proxy_logging.service_logging_obj.async_service_failure_hook.call_args[
|
|
1
|
|
]
|
|
assert call_args["service"] == ServiceTypes.DB
|
|
assert call_args["call_type"] == "failing_function"
|
|
assert call_args["parent_otel_span"] == "test_span"
|
|
assert isinstance(call_args["duration"], float)
|
|
assert isinstance(call_args["start_time"], datetime)
|
|
assert isinstance(call_args["end_time"], datetime)
|
|
assert isinstance(call_args["error"], type(exception))
|
|
else:
|
|
# Assert failure was NOT logged for non-DB errors
|
|
mock_proxy_logging.service_logging_obj.async_service_failure_hook.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_dd_log_db_spend_failure_metrics():
|
|
from litellm._service_logger import ServiceLogging
|
|
from litellm.integrations.datadog.datadog import DataDogLogger
|
|
|
|
dd_logger = DataDogLogger()
|
|
with patch.object(dd_logger, "async_service_failure_hook", new_callable=AsyncMock):
|
|
service_logging_obj = ServiceLogging()
|
|
|
|
litellm.service_callback = [dd_logger]
|
|
|
|
await service_logging_obj.async_service_failure_hook(
|
|
service=ServiceTypes.DB,
|
|
call_type="test_call_type",
|
|
error="test_error",
|
|
duration=1.0,
|
|
)
|