From c727a2ad634e02aca9b28b83ea879e58767e1e63 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Thu, 6 Nov 2025 17:06:03 -0800 Subject: [PATCH] test noma, presidio --- .../guardrails/guardrail_hooks/test_noma.py | 19 +++++++++++++++---- .../guardrail_hooks/test_presidio.py | 6 +++++- 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py index 7a328c03aa4..9a6ab1b8162 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py @@ -267,9 +267,18 @@ class TestNomaGuardrailHooks: # Verify API call details call_args = mock_post.call_args - assert call_args[0][0].endswith("/ai-dr/v2/prompt/scan") - assert call_args[1]["headers"]["X-Noma-AIDR-Application-ID"] == "test-app" - assert call_args[1]["headers"]["Authorization"] == "Bearer test-api-key" + # Verify the URL endpoint + assert call_args.args[0].endswith("/ai-dr/v2/prompt/scan") + # Verify headers and JSON payload + if "headers" in call_args.kwargs: + headers = call_args.kwargs["headers"] + assert "Authorization" in headers + assert headers["Authorization"] == "Bearer test-api-key" + # Verify application ID is in JSON payload (not headers) + if "json" in call_args.kwargs: + json_payload = call_args.kwargs["json"] + assert "x-noma-context" in json_payload + assert json_payload["x-noma-context"]["applicationId"] == "test-app" @pytest.mark.asyncio async def test_pre_call_hook_blocked( @@ -1128,7 +1137,9 @@ class TestNomaImageProcessing: mock_response.json.return_value = noma_response mock_response.raise_for_status = MagicMock() - from litellm.proxy.guardrails.guardrail_hooks.noma.noma import NomaBlockedMessage + from litellm.proxy.guardrails.guardrail_hooks.noma.noma import ( + NomaBlockedMessage, + ) with patch.object( noma_guardrail.async_handler, "post", return_value=mock_response diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py index 234f3a951cd..9543b61ef69 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -577,7 +577,11 @@ async def test_presidio_sets_guardrail_information_in_request_data(): assert "metadata" in request_data assert "standard_logging_guardrail_information" in request_data["metadata"] - guardrail_info = request_data["metadata"]["standard_logging_guardrail_information"] + guardrail_info_list = request_data["metadata"]["standard_logging_guardrail_information"] + assert isinstance(guardrail_info_list, list) + assert len(guardrail_info_list) > 0 + + guardrail_info = guardrail_info_list[0] assert "masked_entity_count" in guardrail_info assert guardrail_info["masked_entity_count"]["EMAIL_ADDRESS"] == 1 assert guardrail_info["masked_entity_count"]["PERSON"] == 1