From b35b269a4639abb9fda2b8b32a51e2007c26d17a Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 26 Nov 2024 19:54:47 -0800
Subject: [PATCH] test_datadog_post_call_failure_hook

---
 litellm/integrations/datadog/datadog.py      |  1 -
 tests/logging_callback_tests/test_datadog.py | 78 ++++++++++++++++++++
 2 files changed, 78 insertions(+), 1 deletion(-)

diff --git a/litellm/integrations/datadog/datadog.py b/litellm/integrations/datadog/datadog.py
index b7348f50e..42d9a38d6 100644
--- a/litellm/integrations/datadog/datadog.py
+++ b/litellm/integrations/datadog/datadog.py
@@ -37,7 +37,6 @@ from litellm.types.integrations.datadog import *
 from litellm.types.services import ServiceLoggerPayload
 from litellm.types.utils import StandardLoggingPayload
 
-from .types import DD_ERRORS, DatadogPayload, DataDogStatus
 from .utils import make_json_serializable
 
 DD_MAX_BATCH_SIZE = 1000  # max number of logs DD API can accept
diff --git a/tests/logging_callback_tests/test_datadog.py b/tests/logging_callback_tests/test_datadog.py
index a93156226..43af1b53b 100644
--- a/tests/logging_callback_tests/test_datadog.py
+++ b/tests/logging_callback_tests/test_datadog.py
@@ -344,3 +344,81 @@ async def test_datadog_logging():
         await asyncio.sleep(5)
     except Exception as e:
         print(e)
+
+
+@pytest.mark.asyncio
+async def test_datadog_post_call_failure_hook():
+    """Test logging proxy failures (e.g., authentication errors) to DataDog"""
+    try:
+        from litellm.integrations.datadog.datadog import DataDogLogger
+
+        os.environ["DD_SITE"] = "https://fake.datadoghq.com"
+        os.environ["DD_API_KEY"] = "anything"
+        dd_logger = DataDogLogger()
+
+        # 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
+
+        # Create a test exception
+        class AuthenticationError(Exception):
+            def __init__(self):
+                self.status_code = 401
+                super().__init__("Invalid API key")
+
+        test_exception = AuthenticationError()
+
+        # Create test request data and user API key dict
+        request_data = {
+            "model": "gpt-4",
+            "messages": [{"role": "user", "content": "Hello"}],
+        }
+
+        user_api_key_dict = UserAPIKeyAuth(
+            api_key="fake_key", user_id="test_user", team_id="test_team"
+        )
+
+        # Call the failure hook
+        await dd_logger.async_post_call_failure_hook(
+            request_data=request_data,
+            original_exception=test_exception,
+            user_api_key_dict=user_api_key_dict,
+        )
+
+        # Wait for the periodic flush
+        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
+
+        # Verify endpoint
+        assert kwargs["url"].endswith("/api/v2/logs"), "Incorrect DataDog endpoint"
+
+        # Decode and verify payload
+        body = kwargs["data"]
+        with gzip.open(io.BytesIO(body), "rb") as f:
+            body = f.read().decode("utf-8")
+
+        body = json.loads(body)
+        assert len(body) == 1, "Expected one log entry"
+
+        log_entry = body[0]
+        assert log_entry["status"] == "error", "Expected error status"
+        assert log_entry["service"] == "litellm-server"
+
+        # Verify message content
+        message = json.loads(log_entry["message"])
+        print("logged message", json.dumps(message, indent=2))
+        assert message["exception"] == "Invalid API key"
+        assert message["error_class"] == "AuthenticationError"
+        assert message["status_code"] == 401
+        assert "traceback" in message
+        assert message["user_api_key_dict"]["api_key"] == "fake_key"
+
+    except Exception as e:
+        pytest.fail(f"Test failed with exception: {str(e)}")