mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
Add more tests
This commit is contained in:
Alex Yaroslavsky
parent
f6f86af312
commit
d4fd4b3dff
1 changed files with 430 additions and 1 deletions
|
|
@ -57,7 +57,7 @@ def test_cato_guard_config():
|
|||
def test_cato_guard_config_no_api_key():
|
||||
litellm.set_verbose = True
|
||||
litellm.guardrail_name_config_map = {}
|
||||
with pytest.raises(CatoNetworksGuardrailMissingSecrets, match="Couldn't get Cato api key"):
|
||||
with pytest.raises(CatoNetworksGuardrailMissingSecrets, match="Couldn't get Cato Networks api key"):
|
||||
init_guardrails_v2(
|
||||
all_guardrails=[
|
||||
{
|
||||
|
|
@ -482,3 +482,432 @@ response_without_detections = Response(
|
|||
status_code=200,
|
||||
request=Request(method="POST", url="http://cato"),
|
||||
)
|
||||
|
||||
|
||||
def _make_response(payload: dict) -> Response:
|
||||
return Response(
|
||||
json=payload,
|
||||
status_code=200,
|
||||
request=Request(method="POST", url="http://cato"),
|
||||
)
|
||||
|
||||
|
||||
def _make_guardrail(api_key: str = "hs-cato-key", **extra) -> CatoNetworksGuardrail:
|
||||
return CatoNetworksGuardrail(api_key=api_key, **extra)
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Constructor coverage
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_init_uses_cato_api_key_env_var(monkeypatch):
|
||||
monkeypatch.setenv("CATO_API_KEY", "from-env")
|
||||
monkeypatch.delenv("CATO_API_BASE", raising=False)
|
||||
guard = CatoNetworksGuardrail()
|
||||
assert guard.api_key == "from-env"
|
||||
assert guard.api_base == "https://api.aisec.catonetworks.com"
|
||||
assert guard.ws_api_base == "wss://api.aisec.catonetworks.com"
|
||||
|
||||
|
||||
def test_init_uses_cato_api_base_env_var(monkeypatch):
|
||||
monkeypatch.setenv("CATO_API_BASE", "https://custom.example.com")
|
||||
guard = _make_guardrail()
|
||||
assert guard.api_base == "https://custom.example.com"
|
||||
assert guard.ws_api_base == "wss://custom.example.com"
|
||||
|
||||
|
||||
def test_init_explicit_args_take_precedence_over_env(monkeypatch):
|
||||
monkeypatch.setenv("CATO_API_KEY", "env-key")
|
||||
monkeypatch.setenv("CATO_API_BASE", "https://env.example.com")
|
||||
guard = CatoNetworksGuardrail(api_key="explicit-key", api_base="https://explicit.example.com")
|
||||
assert guard.api_key == "explicit-key"
|
||||
assert guard.api_base == "https://explicit.example.com"
|
||||
assert guard.ws_api_base == "wss://explicit.example.com"
|
||||
|
||||
|
||||
def test_init_http_api_base_maps_to_ws():
|
||||
guard = _make_guardrail(api_base="http://insecure.example.com")
|
||||
assert guard.ws_api_base == "ws://insecure.example.com"
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# _build_cato_headers direct coverage
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_build_cato_headers_only_required_when_optionals_missing():
|
||||
guard = _make_guardrail()
|
||||
headers = guard._build_cato_headers(
|
||||
hook="pre_call",
|
||||
key_alias=None,
|
||||
user_email=None,
|
||||
litellm_call_id=None,
|
||||
)
|
||||
assert headers["Authorization"] == "Bearer hs-cato-key"
|
||||
assert headers["x-cato-litellm-hook"] == "pre_call"
|
||||
assert "x-cato-litellm-version" in headers
|
||||
assert "x-cato-call-id" not in headers
|
||||
assert "x-cato-user-email" not in headers
|
||||
assert "x-cato-gateway-key-alias" not in headers
|
||||
|
||||
|
||||
def test_build_cato_headers_includes_all_optionals_when_present():
|
||||
guard = _make_guardrail()
|
||||
headers = guard._build_cato_headers(
|
||||
hook="output",
|
||||
key_alias="alias-1",
|
||||
user_email="user@example.com",
|
||||
litellm_call_id="call-123",
|
||||
)
|
||||
assert headers["x-cato-call-id"] == "call-123"
|
||||
assert headers["x-cato-user-email"] == "user@example.com"
|
||||
assert headers["x-cato-gateway-key-alias"] == "alias-1"
|
||||
assert headers["x-cato-litellm-hook"] == "output"
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# call_cato_guardrail (input-side) action branches
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_cato_guardrail_monitor_action_returns_data_unchanged():
|
||||
guard = _make_guardrail()
|
||||
data = {"messages": [{"role": "user", "content": "hi"}]}
|
||||
response = _make_response(
|
||||
{
|
||||
"analysis_result": {"policy_drill_down": {}},
|
||||
"required_action": {"action_type": "monitor_action"},
|
||||
}
|
||||
)
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
||||
return_value=response,
|
||||
):
|
||||
result = await guard.call_cato_guardrail(data, hook="pre_call", key_alias=None)
|
||||
assert result is data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_cato_guardrail_no_required_action_returns_data_unchanged():
|
||||
guard = _make_guardrail()
|
||||
data = {"messages": [{"role": "user", "content": "hi"}]}
|
||||
response = _make_response(
|
||||
{"analysis_result": {"policy_drill_down": {}}, "required_action": None}
|
||||
)
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
||||
return_value=response,
|
||||
):
|
||||
result = await guard.call_cato_guardrail(data, hook="pre_call", key_alias=None)
|
||||
assert result is data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_cato_guardrail_unknown_action_returns_data_unchanged():
|
||||
guard = _make_guardrail()
|
||||
data = {"messages": [{"role": "user", "content": "hi"}]}
|
||||
response = _make_response(
|
||||
{
|
||||
"analysis_result": {"policy_drill_down": {}},
|
||||
"required_action": {"action_type": "totally_made_up"},
|
||||
}
|
||||
)
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
||||
return_value=response,
|
||||
):
|
||||
result = await guard.call_cato_guardrail(data, hook="pre_call", key_alias=None)
|
||||
assert result is data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anonymize_action_without_redacted_chat_returns_data_unchanged():
|
||||
guard = _make_guardrail()
|
||||
data = {"messages": [{"role": "user", "content": "hi"}]}
|
||||
response = _make_response(
|
||||
{
|
||||
"analysis_result": {"policy_drill_down": {}},
|
||||
"required_action": {"action_type": "anonymize_action"},
|
||||
# redacted_chat intentionally absent
|
||||
}
|
||||
)
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
||||
return_value=response,
|
||||
):
|
||||
result = await guard.call_cato_guardrail(data, hook="pre_call", key_alias=None)
|
||||
assert result["messages"] == [{"role": "user", "content": "hi"}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_cato_guardrail_forwards_user_email_from_metadata():
|
||||
guard = _make_guardrail()
|
||||
data = {
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"metadata": {"headers": {"x-cato-user-email": "alice@example.com"}},
|
||||
"litellm_call_id": "call-xyz",
|
||||
}
|
||||
response = _make_response(
|
||||
{"analysis_result": {"policy_drill_down": {}}, "required_action": None}
|
||||
)
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
||||
return_value=response,
|
||||
) as mock_post:
|
||||
await guard.call_cato_guardrail(data, hook="pre_call", key_alias="alias-1")
|
||||
sent_headers = mock_post.call_args.kwargs["headers"]
|
||||
assert sent_headers["x-cato-user-email"] == "alice@example.com"
|
||||
assert sent_headers["x-cato-call-id"] == "call-xyz"
|
||||
assert sent_headers["x-cato-gateway-key-alias"] == "alias-1"
|
||||
assert sent_headers["x-cato-litellm-hook"] == "pre_call"
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Output-side action branches (call_cato_guardrail_on_output / post_call_success_hook)
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_success_hook_block_action_raises():
|
||||
guard = _make_guardrail()
|
||||
request_data = {
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"litellm_call_id": "c-1",
|
||||
}
|
||||
block_response = _make_response(
|
||||
{
|
||||
"analysis_result": {"policy_drill_down": {"PII": {}}},
|
||||
"required_action": {
|
||||
"action_type": "block_action",
|
||||
"detection_message": "blocked output",
|
||||
"policy_name": "PII",
|
||||
},
|
||||
}
|
||||
)
|
||||
llm_response = ModelResponse(
|
||||
choices=[
|
||||
{
|
||||
"finish_reason": "stop",
|
||||
"index": 0,
|
||||
"message": {"content": "secret", "role": "assistant"},
|
||||
}
|
||||
]
|
||||
)
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
||||
return_value=block_response,
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await guard.async_post_call_success_hook(
|
||||
data=request_data,
|
||||
response=llm_response,
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
)
|
||||
assert exc_info.value.status_code == 400
|
||||
assert exc_info.value.detail == "blocked output"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_success_hook_anonymize_action_redacts_content():
|
||||
guard = _make_guardrail()
|
||||
request_data = {"messages": [{"role": "user", "content": "hi"}]}
|
||||
anonymize_response = _make_response(
|
||||
{
|
||||
"analysis_result": {"policy_drill_down": {"PII": {}}},
|
||||
"required_action": {"action_type": "anonymize_action", "policy_name": "PII"},
|
||||
"redacted_chat": {
|
||||
"all_redacted_messages": [
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": "Hello [NAME_1]"},
|
||||
]
|
||||
},
|
||||
}
|
||||
)
|
||||
llm_response = ModelResponse(
|
||||
choices=[
|
||||
{
|
||||
"finish_reason": "stop",
|
||||
"index": 0,
|
||||
"message": {"content": "Hello Brian", "role": "assistant"},
|
||||
}
|
||||
]
|
||||
)
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
||||
return_value=anonymize_response,
|
||||
):
|
||||
result = await guard.async_post_call_success_hook(
|
||||
data=request_data,
|
||||
response=llm_response,
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
)
|
||||
assert isinstance(result, ModelResponse)
|
||||
assert result.choices[0].message.content == "Hello [NAME_1]"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_success_hook_no_action_keeps_content():
|
||||
guard = _make_guardrail()
|
||||
request_data = {"messages": [{"role": "user", "content": "hi"}]}
|
||||
llm_response = ModelResponse(
|
||||
choices=[
|
||||
{
|
||||
"finish_reason": "stop",
|
||||
"index": 0,
|
||||
"message": {"content": "all good", "role": "assistant"},
|
||||
}
|
||||
]
|
||||
)
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
||||
return_value=response_without_detections,
|
||||
):
|
||||
result = await guard.async_post_call_success_hook(
|
||||
data=request_data,
|
||||
response=llm_response,
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
)
|
||||
assert result.choices[0].message.content == "all good"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_success_hook_skips_non_model_response():
|
||||
guard = _make_guardrail()
|
||||
request_data = {"messages": [{"role": "user", "content": "hi"}]}
|
||||
not_a_model_response = {"unexpected": "shape"}
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_post:
|
||||
result = await guard.async_post_call_success_hook(
|
||||
data=request_data,
|
||||
response=not_a_model_response, # type: ignore[arg-type]
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
)
|
||||
mock_post.assert_not_called()
|
||||
assert result is not_a_model_response
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Streaming hook edge cases + forward_the_stream_to_cato serialization
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_stream_unknown_message_terminates_cleanly(monkeypatch):
|
||||
guard = _make_guardrail()
|
||||
|
||||
async def llm_response():
|
||||
if False:
|
||||
yield # pragma: no cover
|
||||
return
|
||||
|
||||
websocket_mock = AsyncMock()
|
||||
messages_from_cato = [b'{"unexpected": "payload"}']
|
||||
websocket_mock.recv = ReceiveMock(messages_from_cato, delay=0.05)
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def connect_mock(*args, **kwargs):
|
||||
yield websocket_mock
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.guardrails.guardrail_hooks.cato_networks.cato_networks.connect",
|
||||
connect_mock,
|
||||
)
|
||||
|
||||
results = []
|
||||
async for chunk in guard.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
response=llm_response(),
|
||||
request_data={"messages": []},
|
||||
):
|
||||
results.append(chunk)
|
||||
assert results == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_stream_forwards_user_email_and_call_id_to_websocket(monkeypatch):
|
||||
guard = _make_guardrail()
|
||||
|
||||
async def llm_response():
|
||||
if False:
|
||||
yield # pragma: no cover
|
||||
return
|
||||
|
||||
websocket_mock = AsyncMock()
|
||||
websocket_mock.recv = ReceiveMock([b'{"done": true}'], delay=0.01)
|
||||
|
||||
captured = {}
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def connect_mock(url, *args, **kwargs):
|
||||
captured["url"] = url
|
||||
captured["headers"] = kwargs.get("additional_headers")
|
||||
yield websocket_mock
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.guardrails.guardrail_hooks.cato_networks.cato_networks.connect",
|
||||
connect_mock,
|
||||
)
|
||||
|
||||
request_data = {
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"metadata": {"headers": {"x-cato-user-email": "bob@example.com"}},
|
||||
"litellm_call_id": "call-stream",
|
||||
}
|
||||
|
||||
async for _ in guard.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(key_alias="kalias"),
|
||||
response=llm_response(),
|
||||
request_data=request_data,
|
||||
):
|
||||
pass
|
||||
|
||||
assert captured["url"].startswith("wss://")
|
||||
assert captured["url"].endswith("/fw/v1/analyze/stream")
|
||||
headers = captured["headers"]
|
||||
assert headers["x-cato-user-email"] == "bob@example.com"
|
||||
assert headers["x-cato-call-id"] == "call-stream"
|
||||
assert headers["x-cato-gateway-key-alias"] == "kalias"
|
||||
assert headers["x-cato-litellm-hook"] == "output"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_forward_the_stream_serializes_chunk_types():
|
||||
guard = _make_guardrail()
|
||||
|
||||
async def chunks():
|
||||
yield ModelResponseStream(id="abc", object="chat.completion.chunk", created=1, model="m", choices=[])
|
||||
yield {"role": "assistant", "content": "hi"}
|
||||
yield "raw-string-chunk"
|
||||
|
||||
websocket_mock = AsyncMock()
|
||||
await guard.forward_the_stream_to_cato(websocket_mock, chunks())
|
||||
|
||||
sent = [c.args[0] for c in websocket_mock.send.mock_calls]
|
||||
assert len(sent) == 4 # 3 chunks + final done sentinel
|
||||
|
||||
# ModelResponseStream → JSON string with expected fields
|
||||
first = json.loads(sent[0])
|
||||
assert first.get("id") == "abc"
|
||||
# dict → json.dumps
|
||||
assert json.loads(sent[1]) == {"role": "assistant", "content": "hi"}
|
||||
# str passthrough
|
||||
assert sent[2] == "raw-string-chunk"
|
||||
# done sentinel
|
||||
assert json.loads(sent[3]) == {"done": True}
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# get_config_model
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_get_config_model_returns_pydantic_class():
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.cato_networks import (
|
||||
CatoNetworksGuardrailConfigModel,
|
||||
)
|
||||
|
||||
assert CatoNetworksGuardrail.get_config_model() is CatoNetworksGuardrailConfigModel
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue