scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

test(proxy): group Bedrock spend tests and patch async_moderation_hook via patch.object

Browse source 
Made-with: Cursor
This commit is contained in:
Milan 2026-04-16 00:44:30 +03:00
parent 19f34b4b9d
commit 36afe10019
No known key found for this signature in database
1 changed files with 28 additions and 23 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

51
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py
View file

@ -6,10 +6,11 @@ import os
import sys
from unittest.mock import AsyncMock, MagicMock, patch
import litellm
import pytest
from fastapi import HTTPException
import litellm
sys.path.insert(0, os.path.abspath("../../../../../.."))
from litellm.caching.caching import DualCache
@ -1195,19 +1196,10 @@ async def test_bedrock_guardrail_blocked_content_with_masking_enabled():
# ---------------------------------------------------------------------------
# L3: _extract_blocked_assessments + _get_http_exception_for_blocked_guardrail
# Regression coverage for case 2026-04-10-internal-bedrock-guardrail-streaming-error.
# Spend logs: guardrail_mode (pre/during/post) vs Bedrock INPUT/OUTPUT
# ---------------------------------------------------------------------------
def _make_guardrail() -> BedrockGuardrail:
return BedrockGuardrail(
guardrail_name="bedrock-pii-guard",
guardrailIdentifier="amgllac6xf3r",
guardrailVersion="1",
)
def test_bedrock_guardrail_uses_native_during_call_hook():
"""during_call must use async_moderation_hook, not unified apply_guardrail(input=request)."""
assert BedrockGuardrail.use_native_during_call_hook is True
@ -1282,27 +1274,40 @@ async def test_during_call_hook_invokes_bedrock_async_moderation_hook():
default_on=True,
)
mock_mod = AsyncMock(return_value=None)
guardrail.async_moderation_hook = mock_mod # type: ignore[method-assign]
original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
try:
litellm.callbacks = [guardrail]
await proxy_logging.during_call_hook(
data={
"model": "gpt-4",
"messages": [{"role": "user", "content": "test"}],
},
user_api_key_dict=UserAPIKeyAuth(
api_key="test_key", user_id="test_user"
),
call_type="completion",
)
with patch.object(guardrail, "async_moderation_hook", new=mock_mod):
await proxy_logging.during_call_hook(
data={
"model": "gpt-4",
"messages": [{"role": "user", "content": "test"}],
},
user_api_key_dict=UserAPIKeyAuth(
api_key="test_key", user_id="test_user"
),
call_type="completion",
)
finally:
litellm.callbacks = original_callbacks
mock_mod.assert_awaited_once()
# ---------------------------------------------------------------------------
# L3: _extract_blocked_assessments + _get_http_exception_for_blocked_guardrail
# Regression coverage for case 2026-04-10-internal-bedrock-guardrail-streaming-error.
# ---------------------------------------------------------------------------
def _make_guardrail() -> BedrockGuardrail:
return BedrockGuardrail(
guardrail_name="bedrock-pii-guard",
guardrailIdentifier="amgllac6xf3r",
guardrailVersion="1",
)
def test_extract_blocked_assessments_pii_entity():
"""L3: PII entity match (BLOCKED) is surfaced with category, type, and matched term."""
g = _make_guardrail()