added test for updating in memory guardrails

Author: ARajan1084

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
+    )