added test for updating in memory guardrails

This commit is contained in:
= 2025-09-17 17:16:49 -07:00 committed by Krrish Dholakia
parent 36e01b7881
commit 40fa415c89

View file

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