scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

test: reorganize unified guardrail tests into nested classes

Browse source
This commit is contained in:
Yuta Saito 2026-01-14 13:38:43 +09:00
parent c5ced033c9
commit 82ed6283fa
1 changed files with 69 additions and 64 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

133
tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py
View file

@ -58,71 +58,76 @@ def _inject_mcp_handler_mapping():
unified_module.endpoint_guardrail_translation_mappings = None
@pytest.mark.asyncio
async def test_pre_call_hook_uses_mcp_event_type():
"""pre_call hook should swap to GuardrailEventHooks.pre_mcp_call for MCP calls."""
handler = UnifiedLLMGuardrails()
guardrail = RecordingGuardrail()
cache = DualCache()
class TestUnifiedLLMGuardrails:
class TestAsyncPreCallHook:
@pytest.mark.asyncio
async def test_uses_mcp_event_type(self):
"""pre_call hook should swap to GuardrailEventHooks.pre_mcp_call for MCP calls."""
handler = UnifiedLLMGuardrails()
guardrail = RecordingGuardrail()
cache = DualCache()
data = {
"guardrail_to_apply": guardrail,
"messages": [{"role": "user", "content": "Tool: test\nArguments: {}"}],
"model": "mcp-tool-call",
}
await handler.async_pre_call_hook(
user_api_key_dict=None,
cache=cache,
data=data,
call_type=CallTypes.call_mcp_tool.value,
)
assert guardrail.event_history == [GuardrailEventHooks.pre_mcp_call]
@pytest.mark.asyncio
async def test_moderation_hook_uses_mcp_event_type():
"""moderation hook should request GuardrailEventHooks.during_mcp_call for MCP calls."""
handler = UnifiedLLMGuardrails()
guardrail = RecordingGuardrail()
data = {
"guardrail_to_apply": guardrail,
"messages": [{"role": "user", "content": "Tool: test\nArguments: {}"}],
"model": "mcp-tool-call",
}
await handler.async_moderation_hook(
data=data,
user_api_key_dict=None,
call_type=CallTypes.call_mcp_tool.value,
)
assert guardrail.event_history == [GuardrailEventHooks.during_mcp_call]
@pytest.mark.asyncio
async def test_moderation_hook_runs_for_anthropic_messages():
"""Ensure anthropic_messages requests still trigger guardrail moderation."""
handler = UnifiedLLMGuardrails()
guardrail = RecordingGuardrail()
data = {
"guardrail_to_apply": guardrail,
"messages": [
{
"role": "user",
"content": "Hello Anthropics",
data = {
"guardrail_to_apply": guardrail,
"messages": [
{"role": "user", "content": "Tool: test\nArguments: {}"}
],
"model": "mcp-tool-call",
}
],
"model": "anthropic.claude-3",
}
await handler.async_moderation_hook(
data=data,
user_api_key_dict=None,
call_type=CallTypes.anthropic_messages.value,
)
await handler.async_pre_call_hook(
user_api_key_dict=None,
cache=cache,
data=data,
call_type=CallTypes.call_mcp_tool.value,
)
assert guardrail.event_history == [GuardrailEventHooks.during_call]
assert guardrail.event_history == [GuardrailEventHooks.pre_mcp_call]
class TestAsyncModerationHook:
@pytest.mark.asyncio
async def test_uses_mcp_event_type(self):
"""moderation hook should request GuardrailEventHooks.during_mcp_call for MCP calls."""
handler = UnifiedLLMGuardrails()
guardrail = RecordingGuardrail()
data = {
"guardrail_to_apply": guardrail,
"messages": [
{"role": "user", "content": "Tool: test\nArguments: {}"}
],
"model": "mcp-tool-call",
}
await handler.async_moderation_hook(
data=data,
user_api_key_dict=None,
call_type=CallTypes.call_mcp_tool.value,
)
assert guardrail.event_history == [GuardrailEventHooks.during_mcp_call]
@pytest.mark.asyncio
async def test_runs_for_anthropic_messages(self):
"""Ensure anthropic_messages requests still trigger guardrail moderation."""
handler = UnifiedLLMGuardrails()
guardrail = RecordingGuardrail()
data = {
"guardrail_to_apply": guardrail,
"messages": [
{
"role": "user",
"content": "Hello Anthropics",
}
],
"model": "anthropic.claude-3",
}
await handler.async_moderation_hook(
data=data,
user_api_key_dict=None,
call_type=CallTypes.anthropic_messages.value,
)
assert guardrail.event_history == [GuardrailEventHooks.during_call]