From 40fa415c8913b48bede0075be1bae4eeafa7ab68 Mon Sep 17 00:00:00 2001
From: = <achintya.rajan03@gmail.com>
Date: Wed, 17 Sep 2025 17:16:49 -0700
Subject: [PATCH] added test for updating in memory guardrails

---
 .../guardrails/test_guardrail_registry.py     | 33 +++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py
index f05ec653fbd..23432b18ca0 100644
--- a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py
+++ b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py
@@ -1,6 +1,9 @@
+from litellm.integrations.custom_guardrail import CustomGuardrail
 from litellm.proxy.guardrails.guardrail_registry import (
     get_guardrail_initializer_from_hooks,
+    InMemoryGuardrailHandler,
 )
+from litellm.types.guardrails import GuardrailEventHooks, Guardrail, LitellmParams
 
 
 def test_get_guardrail_initializer_from_hooks():
@@ -15,3 +18,33 @@ def test_guardrail_class_registry():
     print(f"guardrail_class_registry: {guardrail_class_registry}")
     assert "aim" in guardrail_class_registry
     assert "aporia" in guardrail_class_registry
+
+
+def test_update_in_memory_guardrail():
+    handler = InMemoryGuardrailHandler()
+    handler.guardrail_id_to_custom_guardrail["123"] = CustomGuardrail(
+        guardrail_name="test-guardrail",
+        default_on=False,
+        event_hook=GuardrailEventHooks.pre_call,
+    )
+
+    handler.update_in_memory_guardrail(
+        "123",
+        Guardrail(
+            guardrail_name="test-guardrail",
+            litellm_params=LitellmParams(
+                guardrail="test-guardrail", mode="pre_call", default_on=True
+            ),
+        ),
+    )
+
+    assert (
+        handler.guardrail_id_to_custom_guardrail["123"].should_run_guardrail(
+            data={}, event_type=GuardrailEventHooks.pre_call
+        )
+        is True
+    )
+    assert (
+        handler.guardrail_id_to_custom_guardrail["123"].event_hook
+        is GuardrailEventHooks.pre_call
+    )