scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-05 02:41:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 49

test(guardrails): harden bedrock history assertion

Browse source
This commit is contained in:
CHX606 2026-09-20 22:01:59 +08:00
parent ec5931df5b
commit 853926a84b
1 changed files with 7 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
tests/test_litellm/llms/bedrock/passthrough/guardrail_translation/test_handler.py
View file

@ -1,3 +1,5 @@
from copy import deepcopy
from typing import Final
from unittest.mock import AsyncMock, MagicMock
import pytest
@ -9,7 +11,7 @@ from litellm.llms.bedrock.passthrough.guardrail_translation.handler import (
@pytest.mark.asyncio
async def test_skip_assistant_preserves_converse_history_and_masks_user() -> None:
assistant = {
assistant: Final = {
"role": "assistant",
"content": [
{"text": "old reply"},
@ -22,7 +24,7 @@ async def test_skip_assistant_preserves_converse_history_and_masks_user() -> Non
},
],
}
data = {
data: Final = {
"endpoint": "model/test/converse",
"data": {
"messages": [
@ -31,7 +33,8 @@ async def test_skip_assistant_preserves_converse_history_and_masks_user() -> Non
]
},
}
guardrail = MagicMock()
original_assistant: Final = deepcopy(assistant)
guardrail: Final = MagicMock()
guardrail.apply_guardrail = AsyncMock(return_value={"texts": ["[MASKED]"]})
guardrail.skip_system_message_in_guardrail = False
guardrail.skip_tool_message_in_guardrail = False
@ -40,5 +43,5 @@ async def test_skip_assistant_preserves_converse_history_and_masks_user() -> Non
await BedrockPassthroughGuardrailHandler().process_input_messages(data, guardrail)
assert guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] == ["private"]
assert data["data"]["messages"][0] == assistant
assert data["data"]["messages"][0] == original_assistant
assert data["data"]["messages"][1]["content"][0]["text"] == "[MASKED]"