From d4fd4b3dff20ba44d2e83a2d0c1def6c62526e0d Mon Sep 17 00:00:00 2001
From: Alex Yaroslavsky <alex.yaroslavsky@catonetworks.com>
Date: Tue, 28 Apr 2026 10:53:39 +0300
Subject: [PATCH] Add more tests

---
 .../test_cato_networks_guardrails.py          | 431 +++++++++++++++++-
 1 file changed, 430 insertions(+), 1 deletion(-)

diff --git a/tests/local_testing/test_cato_networks_guardrails.py b/tests/local_testing/test_cato_networks_guardrails.py
index 1805e466ef9..9b5a120950f 100644
--- a/tests/local_testing/test_cato_networks_guardrails.py
+++ b/tests/local_testing/test_cato_networks_guardrails.py
@@ -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