import asyncio import gzip import io import json import logging import os from datetime import datetime as datetime_class from unittest.mock import AsyncMock import pytest import litellm from litellm._logging import verbose_logger from litellm.integrations.datadog.datadog import * from litellm.types.utils import ( StandardLoggingHiddenParams, StandardLoggingMetadata, StandardLoggingModelInformation, StandardLoggingPayload, ) verbose_logger.setLevel(logging.DEBUG) def create_standard_logging_payload() -> StandardLoggingPayload: return StandardLoggingPayload( id="test_id", call_type="completion", response_cost=0.1, response_cost_failure_debug_info=None, status="success", total_tokens=30, prompt_tokens=20, completion_tokens=10, startTime=1234567890.0, endTime=1234567891.0, completionStartTime=1234567890.5, model_map_information=StandardLoggingModelInformation( model_map_key="gpt-4.1-mini", model_map_value=None ), model="gpt-4.1-mini", model_id="model-123", model_group="openai-gpt", api_base="https://api.openai.com", metadata=StandardLoggingMetadata( user_api_key_hash="test_hash", user_api_key_org_id=None, user_api_key_alias="test_alias", user_api_key_team_id="test_team", user_api_key_user_id="test_user", user_api_key_team_alias="test_team_alias", spend_logs_metadata=None, requester_ip_address="127.0.0.1", requester_metadata=None, ), cache_hit=False, cache_key=None, saved_cache_cost=0.0, request_tags=[], end_user=None, requester_ip_address="127.0.0.1", messages=[{"role": "user", "content": "Hello, world!"}], response={"choices": [{"message": {"content": "Hi there!"}}]}, error_str=None, model_parameters={"stream": True}, hidden_params=StandardLoggingHiddenParams( model_id="model-123", cache_key=None, api_base="https://api.openai.com", response_cost="0.1", additional_headers=None, ), ) @pytest.mark.asyncio async def test_create_datadog_logging_payload(): """Test creating a DataDog logging payload from a standard logging object""" dd_logger = DataDogLogger() standard_payload = create_standard_logging_payload() # Create mock kwargs with the standard logging object kwargs = {"standard_logging_object": standard_payload} # Test payload creation dd_payload = dd_logger.create_datadog_logging_payload( kwargs=kwargs, response_obj=None, start_time=datetime_class.now(), end_time=datetime_class.now(), ) # Verify payload structure assert dd_payload["ddsource"] == os.getenv("DD_SOURCE", "litellm") assert dd_payload["service"] == "litellm-server" assert dd_payload["status"] == DataDogStatus.INFO # verify the message field == standard_payload dict_payload = json.loads(dd_payload["message"]) assert dict_payload == standard_payload @pytest.mark.asyncio async def test_datadog_failure_logging(): """Test logging a failure event to DataDog""" dd_logger = DataDogLogger() standard_payload = create_standard_logging_payload() standard_payload["status"] = "failure" # Set status to failure standard_payload["error_str"] = "Test error" kwargs = {"standard_logging_object": standard_payload} dd_payload = dd_logger.create_datadog_logging_payload( kwargs=kwargs, response_obj=None, start_time=datetime_class.now(), end_time=datetime_class.now(), ) assert ( dd_payload["status"] == DataDogStatus.ERROR ) # Verify failure maps to warning status # verify the message field == standard_payload dict_payload = json.loads(dd_payload["message"]) assert dict_payload == standard_payload # verify error_str is in the message field assert "error_str" in dict_payload assert dict_payload["error_str"] == "Test error" @pytest.mark.asyncio async def test_datadog_log_redis_failures(): """ Test that poorly configured Redis is logged as Warning on DataDog """ try: from litellm.caching.caching import Cache from litellm.integrations.datadog.datadog import DataDogLogger litellm.cache = Cache( type="redis", host="badhost", port="6379", password="badpassword" ) os.environ["DD_SITE"] = "https://fake.datadoghq.com" os.environ["DD_API_KEY"] = "anything" dd_logger = DataDogLogger() litellm.callbacks = [dd_logger] litellm.service_callback = ["datadog"] litellm.set_verbose = True # Create a mock for the async_client's post method mock_post = AsyncMock() mock_post.return_value.status_code = 202 mock_post.return_value.text = "Accepted" dd_logger.async_client.post = mock_post # Make the completion call for _ in range(3): response = await litellm.acompletion( model="gpt-4.1-mini", messages=[{"role": "user", "content": "what llm are u"}], max_tokens=10, temperature=0.2, mock_response="Accepted", ) print(response) # Wait for 5 seconds await asyncio.sleep(6) # Assert that the mock was called assert mock_post.called, "HTTP request was not made" # Get the arguments of the last call args, kwargs = mock_post.call_args print("CAll args and kwargs", args, kwargs) # For example, checking if the URL is correct assert kwargs["url"].endswith("/api/v2/logs"), "Incorrect DataDog endpoint" body = kwargs["data"] # use gzip to unzip the body with gzip.open(io.BytesIO(body), "rb") as f: body = f.read().decode("utf-8") print(body) # body is string parse it to dict body = json.loads(body) print(body) failure_events = [log for log in body if log["status"] == "warning"] assert len(failure_events) > 0, "No failure events logged" print("ALL FAILURE/WARN EVENTS", failure_events) for event in failure_events: message = json.loads(event["message"]) assert ( event["status"] == "warning" ), f"Event status is not 'warning': {event['status']}" assert ( message["service"] == "redis" ), f"Service is not 'redis': {message['service']}" assert "error" in message, "No 'error' field in the message" assert message["error"], "Error field is empty" except Exception as e: pytest.fail(f"Test failed with exception: {str(e)}")