From 459ba986074d5f2e12e9933ecb9d9fd35c0b47c6 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 27 Nov 2024 09:40:35 -0800
Subject: [PATCH] test_disable_error_logs

---
 .../test_unit_test_proxy_hooks.py             | 74 +++++++++++++++++++
 1 file changed, 74 insertions(+)
 create mode 100644 tests/proxy_unit_tests/test_unit_test_proxy_hooks.py

diff --git a/tests/proxy_unit_tests/test_unit_test_proxy_hooks.py b/tests/proxy_unit_tests/test_unit_test_proxy_hooks.py
new file mode 100644
index 000000000..5573ad096
--- /dev/null
+++ b/tests/proxy_unit_tests/test_unit_test_proxy_hooks.py
@@ -0,0 +1,74 @@
+import asyncio
+import os
+import sys
+from unittest.mock import Mock, patch, AsyncMock
+import pytest
+from fastapi import Request
+from litellm.proxy.utils import _get_redoc_url, _get_docs_url
+
+sys.path.insert(0, os.path.abspath("../.."))
+import litellm
+
+
+@pytest.mark.asyncio
+async def test_disable_error_logs():
+    """
+    Test that the error logs are not written to the database when disable_error_logs is True
+    """
+    # Mock the necessary components
+    mock_prisma_client = AsyncMock()
+    mock_general_settings = {"disable_error_logs": True}
+
+    with patch(
+        "litellm.proxy.proxy_server.general_settings", mock_general_settings
+    ), patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client):
+
+        # Create a test exception
+        test_exception = Exception("Test error")
+        test_kwargs = {
+            "model": "gpt-4",
+            "exception": test_exception,
+            "optional_params": {},
+            "litellm_params": {"metadata": {}},
+        }
+
+        # Call the failure handler
+        from litellm.proxy.proxy_server import _PROXY_failure_handler
+
+        await _PROXY_failure_handler(
+            kwargs=test_kwargs,
+            completion_response=None,
+            start_time="2024-01-01",
+            end_time="2024-01-01",
+        )
+
+        # Verify prisma client was not called to create error logs
+        if hasattr(mock_prisma_client, "db"):
+            assert not mock_prisma_client.db.litellm_errorlogs.create.called
+
+
+@pytest.mark.asyncio
+async def test_disable_spend_logs():
+    """
+    Test that the spend logs are not written to the database when disable_spend_logs is True
+    """
+    # Mock the necessary components
+    mock_prisma_client = Mock()
+    mock_prisma_client.spend_log_transactions = []
+
+    with patch("litellm.proxy.proxy_server.disable_spend_logs", True), patch(
+        "litellm.proxy.proxy_server.prisma_client", mock_prisma_client
+    ):
+        from litellm.proxy.proxy_server import update_database
+
+        # Call update_database with disable_spend_logs=True
+        await update_database(
+            token="fake-token",
+            response_cost=0.1,
+            user_id="user123",
+            completion_response=None,
+            start_time="2024-01-01",
+            end_time="2024-01-01",
+        )
+        # Verify no spend logs were added
+        assert len(mock_prisma_client.spend_log_transactions) == 0