From 333bd0a9621eeced0e449353fb81f4371d0095bb Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Wed, 6 Nov 2024 00:19:10 +0530
Subject: [PATCH] test(test_proxy_utils.py): add testing for db config update
 logic

---
 litellm/proxy/proxy_server.py           |  5 ---
 tests/local_testing/test_proxy_utils.py | 47 ++++++++++++++++++++++++-
 2 files changed, 46 insertions(+), 6 deletions(-)

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 9f6579242..483abe5ee 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -2392,11 +2392,6 @@ class ProxyConfig:
         for response in responses:
             if response is not None:
                 param_name = getattr(response, "param_name", None)
-                verbose_proxy_logger.info(f"loading {param_name} settings from db")
-                if param_name == "litellm_settings":
-                    verbose_proxy_logger.info(
-                        f"litellm_settings: {response.param_value}"
-                    )
                 param_value = getattr(response, "param_value", None)
                 if param_name is not None and param_value is not None:
                     # check if param_name is already in the config
diff --git a/tests/local_testing/test_proxy_utils.py b/tests/local_testing/test_proxy_utils.py
index 74ef75392..f3f33bad6 100644
--- a/tests/local_testing/test_proxy_utils.py
+++ b/tests/local_testing/test_proxy_utils.py
@@ -10,7 +10,7 @@ sys.path.insert(
     0, os.path.abspath("../..")
 )  # Adds the parent directory to the system path
 import litellm
-
+from unittest.mock import MagicMock, patch, AsyncMock
 
 from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
 from litellm.proxy.auth.auth_utils import is_request_body_safe
@@ -465,3 +465,48 @@ def test_update_internal_user_params():
         updated_data_json["budget_duration"]
         == litellm.default_internal_user_params["budget_duration"]
     )
+
+
+@pytest.mark.asyncio
+async def test_proxy_config_update_from_db():
+    from litellm.proxy.proxy_server import ProxyConfig
+    from pydantic import BaseModel
+
+    proxy_config = ProxyConfig()
+
+    pc = AsyncMock()
+
+    test_config = {
+        "litellm_settings": {
+            "callbacks": ["prometheus", "otel"],
+        }
+    }
+
+    class ReturnValue(BaseModel):
+        param_name: str
+        param_value: dict
+
+    with patch.object(
+        pc,
+        "get_generic_data",
+        new=AsyncMock(
+            return_value=ReturnValue(
+                param_name="litellm_settings",
+                param_value={
+                    "success_callback": "langfuse",
+                },
+            )
+        ),
+    ):
+        new_config = await proxy_config._update_config_from_db(
+            prisma_client=pc,
+            config=test_config,
+            store_model_in_db=True,
+        )
+
+        assert new_config == {
+            "litellm_settings": {
+                "callbacks": ["prometheus", "otel"],
+                "success_callback": "langfuse",
+            }
+        }