litellm/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py
mateo-berri 9dabd72f2d refactor(repositories): type prisma table access with one generic protocol
Every repository handed its `.table` back untyped, so a dozen modules had
each grown a private `_PrismaTableActions` Protocol to paper over it. They
had drifted: some declared `update` as returning the row, others the row or
None, and none agreed on whether `find_many` was covariant

Replace all of them with a single `TableActions[RowT_co]` in
`litellm/repositories/prisma_protocols.py`, keyed to the prisma row each
repository is bound to. Query inputs stay `Mapping[str, object]` so callers
keep passing plain dicts, and `find_many` returns `Sequence` so the row type
stays covariant

Typing the nullable returns honestly surfaced paths that were already
crashing. A team admin could never edit or delete a memory entry owned by
their team: the write-auth check fed a raw prisma row to a helper that
expects the domain model, so `members_with_roles` arrived as plain dicts and
the request died as a 500 instead of applying the edit. Non-admin members hit
the same 500 in place of the 403 they were owed, so refusal and breakage were
indistinguishable. `/v2/model/info?user_models_only=true` dereferenced a
missing user row rather than returning the 400 the route already had, three
team routes dereferenced a team deleted between the read and the write, and
the agent registry dereferenced a missing agent instead of naming it

basedpyright drops 2,132 errors, 1,454 of them reportAny and 73
reportExplicitAny. The dashboard's generated types pick up `string[]` where
they had `unknown[]` for a team's members, admins and models
2026-08-25 12:14:17 +00:00

681 lines
26 KiB
Python

from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy.guardrails.guardrail_registry import (
get_guardrail_initializer_from_hooks,
GuardrailRegistry,
InMemoryGuardrailHandler,
)
from litellm.types.guardrails import GuardrailEventHooks, Guardrail, LitellmParams
def test_get_guardrail_initializer_from_hooks():
initializers = get_guardrail_initializer_from_hooks()
assert "aim" in initializers
def test_guardrail_class_registry():
from litellm.proxy.guardrails.guardrail_registry import guardrail_class_registry
assert "aim" in guardrail_class_registry
assert "aporia" in guardrail_class_registry
def test_noma_registry_resolution():
from litellm.proxy.guardrails.guardrail_hooks.noma.noma import NomaGuardrail
from litellm.proxy.guardrails.guardrail_hooks.noma.noma_v2 import NomaV2Guardrail
from litellm.proxy.guardrails.guardrail_registry import (
guardrail_class_registry,
guardrail_initializer_registry,
)
assert guardrail_class_registry["noma"] is NomaGuardrail
assert guardrail_class_registry["noma_v2"] is NomaV2Guardrail
assert "noma" in guardrail_initializer_registry
assert "noma_v2" in guardrail_initializer_registry
@pytest.mark.parametrize(
"configured, expected",
[(None, True), (False, False), (True, True)],
)
def test_initialize_guardrail_run_in_parallel_preserves_constructor_default(configured, expected):
"""
A guardrail whose constructor sets run_in_parallel=True must keep that default when
the config omits the key; only an explicit config value may override it. The
previous code wrote bool(None)==False on every instance, silently disabling the
opt-in for such guardrails.
"""
from litellm.proxy.guardrails import guardrail_registry as registry_module
def _initializer(litellm_params, guardrail):
return CustomGuardrail(
guardrail_name=guardrail["guardrail_name"],
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
run_in_parallel=True,
)
registry_module.guardrail_initializer_registry["parallel_default_test"] = _initializer
try:
params = {"guardrail": "parallel_default_test", "mode": "pre_call"}
if configured is not None:
params["run_in_parallel"] = configured
handler = InMemoryGuardrailHandler()
result = handler.initialize_guardrail(
guardrail={"guardrail_name": "cf-parallel-default", "litellm_params": params},
)
stored = handler.guardrail_id_to_custom_guardrail[result["guardrail_id"]]
assert stored.run_in_parallel is expected
finally:
registry_module.guardrail_initializer_registry.pop("parallel_default_test", None)
def _register_noop_initializer(guardrail_type: str):
from litellm.proxy.guardrails import guardrail_registry as registry_module
def _initializer(litellm_params, guardrail):
return CustomGuardrail(
guardrail_name=guardrail["guardrail_name"],
event_hook=GuardrailEventHooks.pre_call,
default_on=False,
)
registry_module.guardrail_initializer_registry[guardrail_type] = _initializer
return registry_module
def _config_guardrail(name: str, guardrail_type: str, guardrail_id=None) -> dict:
guardrail = {
"guardrail_name": name,
"litellm_params": {"guardrail": guardrail_type, "mode": "pre_call"},
}
if guardrail_id is not None:
guardrail["guardrail_id"] = guardrail_id
return guardrail
def test_config_guardrail_id_is_stable_across_boots():
"""
Config guardrails used to get a fresh uuid4 per process, so ids from a
previous boot (or another replica) 404'd on /guardrails/{id}/info even
though the guardrail was alive.
"""
registry_module = _register_noop_initializer("stable_id_test")
try:
first_boot = InMemoryGuardrailHandler().initialize_guardrail(
guardrail=_config_guardrail("tooling", "stable_id_test")
)
second_boot = InMemoryGuardrailHandler().initialize_guardrail(
guardrail=_config_guardrail("tooling", "stable_id_test")
)
assert first_boot["guardrail_id"] == second_boot["guardrail_id"]
finally:
registry_module.guardrail_initializer_registry.pop("stable_id_test", None)
def test_explicit_config_guardrail_id_wins_over_derived_id():
registry_module = _register_noop_initializer("explicit_id_test")
try:
result = InMemoryGuardrailHandler().initialize_guardrail(
guardrail=_config_guardrail(
"tooling", "explicit_id_test", guardrail_id="my-explicit-id"
)
)
assert result["guardrail_id"] == "my-explicit-id"
finally:
registry_module.guardrail_initializer_registry.pop("explicit_id_test", None)
def test_duplicate_config_guardrail_names_get_distinct_stable_ids():
"""
Duplicate guardrail_name entries are legitimate (load balancing across
deployments); each occurrence must keep its own id, stable across boots.
"""
registry_module = _register_noop_initializer("dup_name_test")
try:
handler = InMemoryGuardrailHandler()
first = handler.initialize_guardrail(
guardrail=_config_guardrail("dup", "dup_name_test")
)
second = handler.initialize_guardrail(
guardrail=_config_guardrail("dup", "dup_name_test")
)
rebooted_handler = InMemoryGuardrailHandler()
rebooted_first = rebooted_handler.initialize_guardrail(
guardrail=_config_guardrail("dup", "dup_name_test")
)
rebooted_second = rebooted_handler.initialize_guardrail(
guardrail=_config_guardrail("dup", "dup_name_test")
)
assert first["guardrail_id"] != second["guardrail_id"]
assert first["guardrail_id"] == rebooted_first["guardrail_id"]
assert second["guardrail_id"] == rebooted_second["guardrail_id"]
assert len(handler.IN_MEMORY_GUARDRAILS) == 2
finally:
registry_module.guardrail_initializer_registry.pop("dup_name_test", None)
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
def _make_guardrail(guardrail_id: str, name: str = "g") -> Guardrail:
return Guardrail(
guardrail_id=guardrail_id,
guardrail_name=name,
litellm_params=LitellmParams(guardrail=name, mode="pre_call", default_on=False),
)
def test_reconcile_db_guardrails_drops_stale_db_entries_only():
"""
The reconcile pass must drop in-memory entries marked source='db' that are
missing from the DB result, and never touch source='config' entries.
Models the multi-pod case where another pod deleted a DB-backed guardrail.
"""
handler = InMemoryGuardrailHandler()
# Two DB-backed entries on this pod (synced from earlier polling cycles)
handler.IN_MEMORY_GUARDRAILS["db-keep"] = _make_guardrail("db-keep")
handler.IN_MEMORY_GUARDRAILS["db-stale"] = _make_guardrail("db-stale")
handler._sources["db-keep"] = "db"
handler._sources["db-stale"] = "db"
# One config-loaded entry that must survive reconciliation
handler.IN_MEMORY_GUARDRAILS["cfg"] = _make_guardrail("cfg")
handler._sources["cfg"] = "config"
# The DB now only contains db-keep — db-stale was deleted on another pod.
removed = handler.reconcile_db_guardrails(db_guardrail_ids={"db-keep"})
assert removed == ["db-stale"]
assert "db-stale" not in handler.IN_MEMORY_GUARDRAILS
assert "db-stale" not in handler._sources
assert "db-keep" in handler.IN_MEMORY_GUARDRAILS
assert "cfg" in handler.IN_MEMORY_GUARDRAILS
assert handler._sources["cfg"] == "config"
def test_reconcile_does_not_drop_config_entries_missing_from_db():
"""A config-only guardrail (no DB row) must never be reconciled away."""
handler = InMemoryGuardrailHandler()
handler.IN_MEMORY_GUARDRAILS["cfg-only"] = _make_guardrail("cfg-only")
handler._sources["cfg-only"] = "config"
removed = handler.reconcile_db_guardrails(db_guardrail_ids=set())
assert removed == []
assert "cfg-only" in handler.IN_MEMORY_GUARDRAILS
def test_get_source_returns_marker_set_at_insert():
handler = InMemoryGuardrailHandler()
handler.IN_MEMORY_GUARDRAILS["a"] = _make_guardrail("a")
handler._sources["a"] = "db"
handler.IN_MEMORY_GUARDRAILS["b"] = _make_guardrail("b")
handler._sources["b"] = "config"
assert handler.get_source("a") == "db"
assert handler.get_source("b") == "config"
assert handler.get_source("missing") is None
def test_delete_in_memory_guardrail_clears_source_marker():
handler = InMemoryGuardrailHandler()
handler.IN_MEMORY_GUARDRAILS["a"] = _make_guardrail("a")
handler._sources["a"] = "db"
handler.delete_in_memory_guardrail("a")
assert "a" not in handler.IN_MEMORY_GUARDRAILS
assert "a" not in handler._sources
assert handler.get_source("a") is None
def test_list_config_guardrails_excludes_db_sourced():
"""LIT-2529: read surfaces union DB rows with config guardrails; db-sourced
in-memory entries would double-count (or resurrect stale ones), so exclude them."""
handler = InMemoryGuardrailHandler()
handler.IN_MEMORY_GUARDRAILS["cfg"] = _make_guardrail("cfg", name="config-one")
handler._sources["cfg"] = "config"
handler.IN_MEMORY_GUARDRAILS["db"] = _make_guardrail("db", name="db-one")
handler._sources["db"] = "db"
config_guardrails = handler.list_config_guardrails()
assert [g["guardrail_id"] for g in config_guardrails] == ["cfg"]
def test_get_config_guardrail_by_id_returns_config_only():
"""LIT-2529: the detail/logs fallback must return config-owned guardrails and
treat a db-sourced (stale) or missing id as a miss."""
handler = InMemoryGuardrailHandler()
handler.IN_MEMORY_GUARDRAILS["cfg"] = _make_guardrail("cfg", name="config-one")
handler._sources["cfg"] = "config"
handler.IN_MEMORY_GUARDRAILS["db"] = _make_guardrail("db", name="db-one")
handler._sources["db"] = "db"
assert handler.get_config_guardrail_by_id("cfg")["guardrail_name"] == "config-one"
assert handler.get_config_guardrail_by_id("db") is None
assert handler.get_config_guardrail_by_id("missing") is None
def test_initialize_guardrail_early_return_updates_source_marker():
"""
When initialize_guardrail is called for a guardrail that already exists
in memory, the early-return path must still honor the caller's source.
Otherwise a racing polling tick that placed a DB entry in memory first
would leave a later config-init call wrongly marked as 'db' (or vice
versa), and the entry would be reconciled with the wrong classification.
"""
handler = InMemoryGuardrailHandler()
# Simulate a polling tick already placing the entry as DB-backed.
handler.IN_MEMORY_GUARDRAILS["collide"] = _make_guardrail("collide", name="bedrock")
handler._sources["collide"] = "db"
# Config init re-visits the same id (e.g., hot-reload, or UUID collision).
g = Guardrail(
guardrail_id="collide",
guardrail_name="bedrock",
litellm_params=LitellmParams(guardrail="bedrock", mode="pre_call", default_on=False),
)
handler.initialize_guardrail(guardrail=g, source="config")
assert handler.get_source("collide") == "config"
# And the symmetric direction: db sync should override an entry left
# marked as 'config' from a stale init path.
handler.initialize_guardrail(guardrail=g, source="db")
assert handler.get_source("collide") == "db"
def test_sync_guardrail_from_db_marks_source_db_when_unchanged():
"""
sync_guardrail_from_db must enforce source='db' even when params are
unchanged, so a config entry whose UUID happens to collide with a later
DB row gets re-tagged correctly.
"""
handler = InMemoryGuardrailHandler()
g = _make_guardrail("collide")
handler.IN_MEMORY_GUARDRAILS["collide"] = g
handler._sources["collide"] = "config"
handler.sync_guardrail_from_db(g)
assert handler.get_source("collide") == "db"
def _db_litellm_params() -> dict:
"""
Shape produced by GuardrailRegistry.get_all_guardrails_from_db: litellm_params
is a raw dict (not a LitellmParams), holding only the keys originally stored,
a non-schema extra key, and plain-string enum values.
"""
return {
"guardrail": "litellm_content_filter",
"mode": "pre_call",
"default_on": True,
"version": 2,
"blocked_words": [{"keyword": "secret", "action": "BLOCK"}],
}
def test_unchanged_db_params_do_not_register_as_changed():
"""
A DB poll returns litellm_params as a raw dict while the in-memory copy is a
LitellmParams whose model_dump() fills every field default and coerces enums.
The two shapes must compare equal when the config is identical; otherwise
every poll cycle re-initializes the guardrail indefinitely.
"""
handler = InMemoryGuardrailHandler()
raw = _db_litellm_params()
gid = "11111111-1111-1111-1111-111111111111"
handler.IN_MEMORY_GUARDRAILS[gid] = Guardrail(
guardrail_id=gid,
guardrail_name="cf",
litellm_params=LitellmParams(**raw),
)
new = Guardrail(guardrail_id=gid, guardrail_name="cf", litellm_params=dict(raw))
assert handler._has_guardrail_params_changed(gid, new) is False
def test_changed_db_params_register_as_changed():
"""Normalizing both sides must still surface a genuine config change."""
handler = InMemoryGuardrailHandler()
raw = _db_litellm_params()
gid = "22222222-2222-2222-2222-222222222222"
handler.IN_MEMORY_GUARDRAILS[gid] = Guardrail(
guardrail_id=gid,
guardrail_name="cf",
litellm_params=LitellmParams(**raw),
)
changed = {**raw, "blocked_words": [{"keyword": "different", "action": "BLOCK"}]}
new = Guardrail(guardrail_id=gid, guardrail_name="cf", litellm_params=changed)
assert handler._has_guardrail_params_changed(gid, new) is True
def test_unnormalizable_db_params_register_as_changed_without_raising():
"""
A DB row whose litellm_params fail LitellmParams validation must not crash the
poll loop. The comparison falls back to treating the guardrail as changed so it
re-initializes (and surfaces the bad row in logs) rather than propagating the
validation error up through the polling cycle.
"""
handler = InMemoryGuardrailHandler()
raw = _db_litellm_params()
gid = "55555555-5555-5555-5555-555555555555"
handler.IN_MEMORY_GUARDRAILS[gid] = Guardrail(
guardrail_id=gid,
guardrail_name="cf",
litellm_params=LitellmParams(**raw),
)
malformed = {**raw, "default_on": "not-a-bool-xyz"}
new = Guardrail(guardrail_id=gid, guardrail_name="cf", litellm_params=malformed)
assert handler._has_guardrail_params_changed(gid, new) is True
def _all_callback_lists():
import litellm
return [
litellm.callbacks,
litellm.success_callback,
litellm.failure_callback,
litellm._async_success_callback,
litellm._async_failure_callback,
]
def test_delete_in_memory_guardrail_removes_callback_from_all_lists():
"""
Request handling promotes guardrail callbacks from litellm.callbacks into the
success/failure/async lists. delete_in_memory_guardrail must purge the callback
from every list, otherwise a re-initialized guardrail leaves its old instance
stranded in those lists and instances accumulate.
"""
handler = InMemoryGuardrailHandler()
callback = CustomGuardrail(
guardrail_name="cf-delete",
default_on=True,
event_hook=GuardrailEventHooks.pre_call,
)
gid = "33333333-3333-3333-3333-333333333333"
handler.IN_MEMORY_GUARDRAILS[gid] = _make_guardrail(gid, "cf-delete")
handler._sources[gid] = "db"
handler.guardrail_id_to_custom_guardrail[gid] = callback
lists = _all_callback_lists()
snapshots = [list(cb_list) for cb_list in lists]
try:
for cb_list in lists:
cb_list.append(callback)
handler.delete_in_memory_guardrail(gid)
for cb_list in lists:
assert callback not in cb_list
finally:
for cb_list, snapshot in zip(lists, snapshots):
cb_list[:] = snapshot
def test_repeated_db_sync_does_not_accumulate_runner_instances():
"""
End-to-end regression for the OOM: across repeated DB polls (with the config
genuinely changing each cycle to force re-initialization), exactly one live
guardrail instance must exist across all callback lists. On the unfixed code
the stale instance lingers in the success/failure lists and the distinct count
climbs above one.
"""
import litellm
handler = InMemoryGuardrailHandler()
gid = "44444444-4444-4444-4444-444444444444"
name = "cf-accum"
def db_guardrail(word: str) -> Guardrail:
params = {
**_db_litellm_params(),
"blocked_words": [{"keyword": word, "action": "BLOCK"}],
}
return Guardrail(guardrail_id=gid, guardrail_name=name, litellm_params=params)
def promote_into_request_lists() -> None:
manager = litellm.logging_callback_manager
for callback in list(litellm.callbacks):
manager.add_litellm_success_callback(callback)
manager.add_litellm_failure_callback(callback)
manager.add_litellm_async_success_callback(callback)
manager.add_litellm_async_failure_callback(callback)
def distinct_runner_instances() -> int:
seen = set()
for callback in litellm.logging_callback_manager._get_all_callbacks():
if isinstance(callback, CustomGuardrail) and getattr(callback, "guardrail_name", None) == name:
seen.add(id(callback))
return len(seen)
lists = _all_callback_lists()
snapshots = [list(cb_list) for cb_list in lists]
try:
for cycle in range(5):
handler.sync_guardrail_from_db(db_guardrail(f"word-{cycle}"))
promote_into_request_lists()
assert distinct_runner_instances() == 1
finally:
for cb_list, snapshot in zip(lists, snapshots):
cb_list[:] = snapshot
def _judge_guardrail(guardrail_id: str) -> Guardrail:
return Guardrail(
guardrail_id=guardrail_id,
guardrail_name="quality-judge",
litellm_params={
"guardrail": "llm_as_a_judge",
"mode": "post_call",
"judge_model": "my-judge-alias",
"overall_threshold": 80,
"on_failure": "log",
"criteria": [{"name": "helpfulness", "weight": 100, "description": "helpful?"}],
},
)
def test_db_synced_judge_guardrail_uses_lazy_router_provider():
"""A judge guardrail created/synced through a DB path must resolve the active
Router lazily at call time (issue: UI-created guardrails failed open because the
Router was captured at construction; a guardrail created before the Router
existed captured None and never recovered). Asserting the default provider is
wired guarantees the instance reads the live global rather than a stale value."""
from litellm.proxy.guardrails.guardrail_hooks.llm_as_a_judge import (
LLMAsAJudgeGuardrail,
_default_router_provider,
)
handler = InMemoryGuardrailHandler()
lists = _all_callback_lists()
snapshots = [list(cb_list) for cb_list in lists]
try:
handler.sync_guardrail_from_db(_judge_guardrail("judge-db"))
instance = handler.guardrail_id_to_custom_guardrail["judge-db"]
assert isinstance(instance, LLMAsAJudgeGuardrail)
assert instance._router_provider is _default_router_provider
finally:
for cb_list, snapshot in zip(lists, snapshots):
cb_list[:] = snapshot
def test_reinitialized_judge_guardrail_uses_lazy_router_provider():
from litellm.proxy.guardrails.guardrail_hooks.llm_as_a_judge import (
LLMAsAJudgeGuardrail,
_default_router_provider,
)
handler = InMemoryGuardrailHandler()
lists = _all_callback_lists()
snapshots = [list(cb_list) for cb_list in lists]
try:
handler.reinitialize_guardrail(_judge_guardrail("judge-reinit"), source="db")
instance = handler.guardrail_id_to_custom_guardrail["judge-reinit"]
assert isinstance(instance, LLMAsAJudgeGuardrail)
assert instance._router_provider is _default_router_provider
finally:
for cb_list, snapshot in zip(lists, snapshots):
cb_list[:] = snapshot
class TestScanOnlyToolResultsInitRefusal:
"""A guardrail whose role filtering never scans tool results must be rejected at
initialization when configured with scan_only_tool_results, instead of booting a
proxy that silently scans nothing on every request."""
def _initialize(self, name: str, params: dict):
lists = _all_callback_lists()
snapshots = [list(cb_list) for cb_list in lists]
try:
return InMemoryGuardrailHandler().initialize_guardrail(
guardrail={"guardrail_name": name, "litellm_params": params},
)
finally:
for cb_list, snapshot in zip(lists, snapshots):
cb_list[:] = snapshot
def test_panw_prisma_airs_with_scan_only_tool_results_is_rejected(self):
with pytest.raises(ValueError, match="never scans tool results"):
self._initialize(
"panw-scan-only-combo",
{
"guardrail": "panw_prisma_airs",
"mode": "pre_call",
"api_key": "test-key",
"profile_name": "test-profile",
"scan_only_tool_results": True,
},
)
def test_bedrock_latest_role_with_scan_only_tool_results_is_rejected(self):
with pytest.raises(ValueError, match="never scans tool results"):
self._initialize(
"bedrock-latest-role-scan-only-combo",
{
"guardrail": "bedrock",
"mode": "pre_call",
"guardrailIdentifier": "gr-1",
"guardrailVersion": "1",
"experimental_use_latest_role_message_only": True,
"scan_only_tool_results": True,
},
)
def test_bedrock_without_latest_role_accepts_scan_only_tool_results(self):
result = self._initialize(
"bedrock-scan-only-ok",
{
"guardrail": "bedrock",
"mode": "pre_call",
"guardrailIdentifier": "gr-1",
"guardrailVersion": "1",
"scan_only_tool_results": True,
},
)
assert result is not None
def test_prompt_security_default_tool_filtering_rejects_scan_only_tool_results(self, monkeypatch):
monkeypatch.delenv("PROMPT_SECURITY_CHECK_TOOL_RESULTS", raising=False)
with pytest.raises(ValueError, match="never scans tool results"):
self._initialize(
"prompt-security-scan-only-combo",
{
"guardrail": "prompt_security",
"mode": "pre_call",
"api_key": "test-key",
"api_base": "https://ps.example.com",
"scan_only_tool_results": True,
},
)
def test_prompt_security_check_tool_results_accepts_scan_only_tool_results(self, monkeypatch):
monkeypatch.setenv("PROMPT_SECURITY_CHECK_TOOL_RESULTS", "true")
result = self._initialize(
"prompt-security-scan-only-ok",
{
"guardrail": "prompt_security",
"mode": "pre_call",
"api_key": "test-key",
"api_base": "https://ps.example.com",
"scan_only_tool_results": True,
},
)
assert result is not None
def test_skip_tool_message_with_scan_only_tool_results_is_rejected(self):
with pytest.raises(ValueError, match="skip_tool_message_in_guardrail are enabled together"):
self._initialize(
"bedrock-skip-tool-scan-only-combo",
{
"guardrail": "bedrock",
"mode": "pre_call",
"guardrailIdentifier": "gr-1",
"guardrailVersion": "1",
"skip_tool_message_in_guardrail": True,
"scan_only_tool_results": True,
},
)
@pytest.mark.asyncio
async def test_update_guardrail_in_db_raises_when_row_missing():
prisma_client = MagicMock()
prisma_client.db.litellm_guardrailstable.update = AsyncMock(return_value=None)
with pytest.raises(
Exception,
match=r"^Error updating guardrail in DB: Guardrail not found, passed guardrail_id=missing-guardrail$",
):
await GuardrailRegistry().update_guardrail_in_db(
guardrail_id="missing-guardrail",
guardrail=Guardrail(
guardrail_name="missing-guardrail",
litellm_params=LitellmParams(guardrail="bedrock", mode="pre_call"),
),
prisma_client=prisma_client,
)