From a552da5eef252f95d922c52e60308f51dd06495b Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Mon, 10 Feb 2025 20:40:56 -0800
Subject: [PATCH] test: update test

---
 litellm/types/guardrails.py                         |  2 +-
 .../logging_callback_tests/test_custom_guardrail.py | 13 ++++++++++++-
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py
index 724b0eefcf9..b7018fe29f9 100644
--- a/litellm/types/guardrails.py
+++ b/litellm/types/guardrails.py
@@ -156,7 +156,7 @@ class GuardrailLiteLLMParamsResponse(BaseModel):
 
 
 class GuardrailInfoResponse(BaseModel):
-    guardrail_name: Optional[str]
+    guardrail_name: str
     litellm_params: GuardrailLiteLLMParamsResponse
     guardrail_info: Optional[Dict]
 
diff --git a/tests/logging_callback_tests/test_custom_guardrail.py b/tests/logging_callback_tests/test_custom_guardrail.py
index ae8b8c0c1dc..af1270756f2 100644
--- a/tests/logging_callback_tests/test_custom_guardrail.py
+++ b/tests/logging_callback_tests/test_custom_guardrail.py
@@ -156,6 +156,12 @@ def test_get_guardrails_list_response():
     sample_config = [
         {
             "guardrail_name": "test-guard",
+            "litellm_params": {
+                "guardrail": "test-guard",
+                "mode": "pre_call",
+                "api_key": "test-api-key",
+                "api_base": "test-api-base",
+            },
             "guardrail_info": {
                 "params": [
                     {
@@ -188,7 +194,12 @@ def test_get_guardrails_list_response():
     assert len(empty_response.guardrails) == 0
 
     # Test case 3: Missing optional fields
-    minimal_config = [{"guardrail_name": "minimal-guard"}]
+    minimal_config = [
+        {
+            "guardrail_name": "minimal-guard",
+            "litellm_params": {"guardrail": "minimal-guard", "mode": "pre_call"},
+        }
+    ]
     minimal_response = _get_guardrails_list_response(minimal_config)
     assert isinstance(minimal_response, ListGuardrailsResponse)
     assert len(minimal_response.guardrails) == 1