340 lines
11 KiB
Python
340 lines
11 KiB
Python
"""
|
|
Unit tests for Zscaler AI Guard guardrail
|
|
|
|
Tests covering:
|
|
- API call handling (ALLOW, BLOCK, errors)
|
|
- Policy ID precedence (metadata > user_api_key > team > init)
|
|
- Boolean config parameter handling (send_user_api_key_*)
|
|
- Metadata resolution (pre-call vs post-call)
|
|
- Header preparation with kwargs
|
|
- resolve-and-execute-policy endpoint (policyId omission)
|
|
"""
|
|
|
|
import pytest
|
|
from unittest.mock import AsyncMock, Mock, patch
|
|
from fastapi import HTTPException
|
|
from litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard import ZscalerAIGuard
|
|
import asyncio
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_make_zscaler_ai_guard_api_call_allow():
|
|
"""Test Zscaler AI Guard API call when response action is 'ALLOW'."""
|
|
# Mock the Zscaler AI Guard API response
|
|
mock_response = Mock()
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {
|
|
"statusCode": 200,
|
|
"action": "ALLOW",
|
|
"zscaler_ai_guard_response": {},
|
|
}
|
|
|
|
guardrail = ZscalerAIGuard(
|
|
api_key="test_api_key", api_base="http://example.com", policy_id=1
|
|
)
|
|
with patch.object(
|
|
guardrail, "_send_request", new_callable=AsyncMock
|
|
) as mock_send_request:
|
|
mock_send_request.return_value = mock_response
|
|
result = await guardrail.make_zscaler_ai_guard_api_call(
|
|
guardrail.zscaler_ai_guard_url,
|
|
guardrail.api_key,
|
|
guardrail.policy_id,
|
|
"IN",
|
|
"Test content",
|
|
)
|
|
|
|
assert result["action"] == "ALLOW"
|
|
assert (
|
|
result["zscaler_ai_guard_response"]["zscaler_ai_guard_response"] == {}
|
|
) # Validating response structure
|
|
assert result["direction"] == "IN" # Check additional fields returned
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_make_zscaler_ai_guard_api_call_block():
|
|
"""Test Zscaler AI Guard API call when response action is 'BLOCK'."""
|
|
# Mock the Zscaler AI Guard API response
|
|
mock_response = Mock()
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {
|
|
"statusCode": 200,
|
|
"action": "BLOCK",
|
|
"transactionId": "12345",
|
|
"detectorResponses": {"detector-1": {"triggered": True, "action": "BLOCK"}},
|
|
}
|
|
|
|
guardrail = ZscalerAIGuard(
|
|
api_key="test_api_key", api_base="http://example.com", policy_id=1
|
|
)
|
|
with patch.object(
|
|
guardrail, "_send_request", new_callable=AsyncMock
|
|
) as mock_send_request:
|
|
mock_send_request.return_value = mock_response
|
|
result = await guardrail.make_zscaler_ai_guard_api_call(
|
|
guardrail.zscaler_ai_guard_url,
|
|
guardrail.api_key,
|
|
guardrail.policy_id,
|
|
"IN",
|
|
"Blocked content",
|
|
)
|
|
|
|
assert result["action"] == "BLOCK"
|
|
assert result["zscaler_ai_guard_response"]["transactionId"] == "12345"
|
|
assert (
|
|
result["zscaler_ai_guard_response"]["detectorResponses"]["detector-1"][
|
|
"action"
|
|
]
|
|
== "BLOCK"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_make_zscaler_ai_guard_api_call_request_exception():
|
|
"""Test Zscaler AI Guard API call where an exception in the request occurs."""
|
|
guardrail = ZscalerAIGuard(
|
|
api_key="test_api_key", api_base="http://example.com", policy_id=1
|
|
)
|
|
with patch.object(
|
|
guardrail, "_send_request", new_callable=AsyncMock
|
|
) as mock_send_request:
|
|
mock_send_request.side_effect = Exception("Connection error")
|
|
|
|
with pytest.raises(HTTPException) as e:
|
|
await guardrail.make_zscaler_ai_guard_api_call(
|
|
guardrail.zscaler_ai_guard_url,
|
|
guardrail.api_key,
|
|
guardrail.policy_id,
|
|
"IN",
|
|
"Error content",
|
|
)
|
|
|
|
assert e.value.status_code == 500
|
|
assert "Connection error" in e.value.detail["reason"]
|
|
|
|
|
|
def test_extract_blocking_info():
|
|
"""Test extract_blocking_info method."""
|
|
guardrail = ZscalerAIGuard(
|
|
api_key="test_api_key", api_base="http://example.com", policy_id=1
|
|
)
|
|
|
|
response = {
|
|
"transactionId": "12345",
|
|
"detectorResponses": {
|
|
"detector1": {"triggered": True, "action": "BLOCK"},
|
|
"detector2": {"triggered": False, "action": "ALLOW"},
|
|
},
|
|
}
|
|
|
|
blocking_info = guardrail.extract_blocking_info(response)
|
|
|
|
assert blocking_info["transactionId"] == "12345"
|
|
assert blocking_info["blockingDetectors"] == ["detector1"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch(
|
|
"litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard.ZscalerAIGuard.make_zscaler_ai_guard_api_call",
|
|
new_callable=AsyncMock,
|
|
)
|
|
async def test_apply_guardrail_text_concatenation(mock_api_call):
|
|
"""
|
|
Test that `apply_guardrail` correctly concatenates texts.
|
|
"""
|
|
guardrail = ZscalerAIGuard(policy_id=100)
|
|
inputs = {"texts": ["Hello", "world"]}
|
|
request_data = {}
|
|
|
|
await guardrail.apply_guardrail(inputs, request_data, "request")
|
|
|
|
mock_api_call.assert_called_once()
|
|
call_args = mock_api_call.call_args
|
|
assert call_args.kwargs["content"] == "Hello world"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch(
|
|
"litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard.ZscalerAIGuard.make_zscaler_ai_guard_api_call",
|
|
new_callable=AsyncMock,
|
|
)
|
|
async def test_policy_id_from_request_metadata(mock_api_call):
|
|
"""
|
|
Test policy_id is picked from request metadata (highest precedence).
|
|
"""
|
|
guardrail = ZscalerAIGuard(policy_id=100)
|
|
inputs = {"texts": ["test"]}
|
|
request_data = {
|
|
"metadata": {
|
|
"zguard_policy_id": 1,
|
|
"user_api_key_metadata": {"zguard_policy_id": 2},
|
|
"team_metadata": {"zguard_policy_id": 3},
|
|
}
|
|
}
|
|
|
|
await guardrail.apply_guardrail(inputs, request_data, "request")
|
|
|
|
mock_api_call.assert_called_once()
|
|
assert mock_api_call.call_args.kwargs["policy_id"] == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch(
|
|
"litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard.ZscalerAIGuard.make_zscaler_ai_guard_api_call",
|
|
new_callable=AsyncMock,
|
|
)
|
|
async def test_policy_id_from_user_api_key_metadata(mock_api_call):
|
|
"""
|
|
Test policy_id is picked from user_api_key_metadata (2nd precedence).
|
|
"""
|
|
guardrail = ZscalerAIGuard(policy_id=100)
|
|
inputs = {"texts": ["test"]}
|
|
request_data = {
|
|
"metadata": {
|
|
"user_api_key_metadata": {"zguard_policy_id": 2},
|
|
"team_metadata": {"zguard_policy_id": 3},
|
|
}
|
|
}
|
|
|
|
await guardrail.apply_guardrail(inputs, request_data, "request")
|
|
|
|
mock_api_call.assert_called_once()
|
|
assert mock_api_call.call_args.kwargs["policy_id"] == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch(
|
|
"litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard.ZscalerAIGuard.make_zscaler_ai_guard_api_call",
|
|
new_callable=AsyncMock,
|
|
)
|
|
async def test_policy_id_from_team_metadata(mock_api_call):
|
|
"""
|
|
Test policy_id is picked from team_metadata (3rd precedence).
|
|
"""
|
|
guardrail = ZscalerAIGuard(policy_id=100)
|
|
inputs = {"texts": ["test"]}
|
|
request_data = {"metadata": {"team_metadata": {"zguard_policy_id": 3}}}
|
|
|
|
await guardrail.apply_guardrail(inputs, request_data, "request")
|
|
|
|
mock_api_call.assert_called_once()
|
|
assert mock_api_call.call_args.kwargs["policy_id"] == 3
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch(
|
|
"litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard.ZscalerAIGuard.make_zscaler_ai_guard_api_call",
|
|
new_callable=AsyncMock,
|
|
)
|
|
async def test_policy_id_from_init(mock_api_call):
|
|
"""
|
|
Test policy_id is picked from guardrail initialization (lowest precedence).
|
|
"""
|
|
guardrail = ZscalerAIGuard(policy_id=100)
|
|
inputs = {"texts": ["test"]}
|
|
request_data = {"metadata": {}}
|
|
|
|
await guardrail.apply_guardrail(inputs, request_data, "request")
|
|
|
|
mock_api_call.assert_called_once()
|
|
assert mock_api_call.call_args.kwargs["policy_id"] == 100
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch(
|
|
"litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard.ZscalerAIGuard.make_zscaler_ai_guard_api_call",
|
|
new_callable=AsyncMock,
|
|
)
|
|
async def test_policy_id_zero_from_request_metadata(mock_api_call):
|
|
"""
|
|
Test policy_id=0 is correctly picked. Make sure pick exact policy_id which users set
|
|
"""
|
|
guardrail = ZscalerAIGuard(policy_id=100)
|
|
inputs = {"texts": ["test"]}
|
|
request_data = {
|
|
"metadata": {
|
|
"zguard_policy_id": 0,
|
|
}
|
|
}
|
|
await guardrail.apply_guardrail(inputs, request_data, "request")
|
|
mock_api_call.assert_called_once()
|
|
assert mock_api_call.call_args.kwargs["policy_id"] == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_use_config_send_user_api_key_alias_when_true():
|
|
"""Test that send_user_api_key_alias=True from config is used (not overridden by env)"""
|
|
guardrail = ZscalerAIGuard(
|
|
api_key="test_key",
|
|
send_user_api_key_alias=True,
|
|
)
|
|
assert guardrail.send_user_api_key_alias is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_preserve_policy_id_zero_in_init():
|
|
"""Test that policy_id=0 is preserved (not treated as falsy and overridden by env)"""
|
|
guardrail = ZscalerAIGuard(
|
|
api_key="test_key",
|
|
policy_id=0,
|
|
)
|
|
assert guardrail.policy_id == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_resolve_from_litellm_metadata_during_post_call():
|
|
"""Test that user_api_key_alias is resolved from litellm_metadata during post-call"""
|
|
request_data = {"litellm_metadata": {"user_api_key_alias": "test-alias-post-call"}}
|
|
result = ZscalerAIGuard._resolve_metadata_value(request_data, "user_api_key_alias")
|
|
assert result == "test-alias-post-call"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_resolve_user_api_key_key_alias_mapping():
|
|
"""Test key_alias -> user_api_key_key_alias mapping in litellm_metadata"""
|
|
request_data = {"litellm_metadata": {"user_api_key_key_alias": "test-key-alias"}}
|
|
result = ZscalerAIGuard._resolve_metadata_value(request_data, "user_api_key_alias")
|
|
assert result == "test-key-alias"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_include_user_api_key_alias_header():
|
|
"""Test that user-api-key-alias header is included when send_user_api_key_alias is True"""
|
|
guardrail = ZscalerAIGuard(
|
|
api_key="test_key",
|
|
send_user_api_key_alias=True,
|
|
)
|
|
headers = guardrail._prepare_headers("test_key", user_api_key_alias="test-alias")
|
|
assert headers.get("user-api-key-alias") == "test-alias"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_omit_policy_id_when_zero_or_negative():
|
|
"""Test that policyId is omitted from request body when policy_id <= 0 (for resolve-and-execute-policy)"""
|
|
guardrail = ZscalerAIGuard(
|
|
api_key="test_key",
|
|
policy_id=-1,
|
|
)
|
|
|
|
mock_response = Mock()
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {
|
|
"statusCode": 200,
|
|
"action": "ALLOW",
|
|
}
|
|
|
|
with patch.object(guardrail, "_send_request", new_callable=AsyncMock) as mock_send:
|
|
mock_send.return_value = mock_response
|
|
|
|
await guardrail.make_zscaler_ai_guard_api_call(
|
|
zscaler_ai_guard_url="http://example.com",
|
|
api_key="test_key",
|
|
policy_id=-1,
|
|
direction="OUT",
|
|
content="test content",
|
|
)
|
|
|
|
call_args = mock_send.call_args
|
|
data = call_args[0][2] # Third positional arg is data
|
|
assert "policyId" not in data
|