From 150363be0b1c2e52c856c4ca63e39e7167476c7b Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Tue, 27 Jan 2026 12:30:47 -0800 Subject: [PATCH] fix(test): mock Lakera API in PII masking test for deterministic behavior Changed from integration test to unit test by mocking the Lakera API response. This makes the test: - Deterministic and not dependent on external API behavior - Consistent with other tests in the file which are all mocked - Able to test the masking logic regardless of Lakera's detection of test data - Faster and more reliable in CI/CD The mock response includes both credit card and email in the payload with proper start/end positions so we can verify the masking logic works correctly. --- tests/guardrails_tests/test_lakera_v2.py | 77 +++++++++++++++--------- 1 file changed, 47 insertions(+), 30 deletions(-) diff --git a/tests/guardrails_tests/test_lakera_v2.py b/tests/guardrails_tests/test_lakera_v2.py index 7d20eb1dfe8..2a8731d5ecd 100644 --- a/tests/guardrails_tests/test_lakera_v2.py +++ b/tests/guardrails_tests/test_lakera_v2.py @@ -22,41 +22,58 @@ async def test_lakera_pre_call_hook_for_pii_masking(): # Setup the guardrail with specific entities config litellm._turn_on_debug() lakera_guardrail = LakeraAIGuardrail( - api_key=os.environ.get("LAKERA_API_KEY"), + api_key="test_key", ) - # Create a sample request with PII data - # Note: Using test email only, as test credit card numbers (like 4111-1111-1111-1111) - # may not be consistently flagged by Lakera's API for masking in payload - data = { - "messages": [ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "My email is test@example.com and my phone number is 555-123-4567"} + # Mock response with PII detections in payload (with start/end positions for masking) + mock_response = { + 'payload': [ + {'detector_type': 'pii/credit_card', 'start': 18, 'end': 37, 'message_id': 1}, # "4111-1111-1111-1111" + {'detector_type': 'pii/email', 'start': 54, 'end': 70, 'message_id': 1}, # "test@example.com" ], - "model": "gpt-3.5-turbo", - "metadata": {} + 'flagged': True, + 'breakdown': [ + {'detector_type': 'pii/credit_card', 'detected': True, 'message_id': 1}, + {'detector_type': 'pii/email', 'detected': True, 'message_id': 1}, + ] } - # Mock objects needed for the pre-call hook - user_api_key_dict = UserAPIKeyAuth(api_key="test_key") - cache = DualCache() - - # Call the pre-call hook with the specified call type - modified_data = await lakera_guardrail.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=cache, - data=data, - call_type="completion" - ) - print(modified_data) - - # Verify the messages have been modified to mask PII - assert modified_data["messages"][0]["content"] == "You are a helpful assistant." # System prompt should be unchanged - - user_message = modified_data["messages"][1]["content"] - # Verify email is masked (Lakera should return this in payload for masking) - assert "test@example.com" not in user_message - assert "[MASKED EMAIL]" in user_message or "****" in user_message # Accept either masking format + with patch.object(lakera_guardrail, 'call_v2_guard', new_callable=AsyncMock) as mock_call: + mock_call.return_value = (mock_response, {}) + + # Create a sample request with PII data + data = { + "messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "My credit card is 4111-1111-1111-1111 and my email is test@example.com. My phone number is 555-123-4567"} + ], + "model": "gpt-3.5-turbo", + "metadata": {} + } + + # Mock objects needed for the pre-call hook + user_api_key_dict = UserAPIKeyAuth(api_key="test_key") + cache = DualCache() + + # Call the pre-call hook with the specified call type + modified_data = await lakera_guardrail.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="completion" + ) + print(modified_data) + + # Verify the messages have been modified to mask PII + assert modified_data["messages"][0]["content"] == "You are a helpful assistant." # System prompt should be unchanged + + user_message = modified_data["messages"][1]["content"] + # Verify both credit card and email are masked + assert "4111-1111-1111-1111" not in user_message + assert "test@example.com" not in user_message + # Verify masking placeholders are present + assert "[MASKED CREDIT_CARD]" in user_message + assert "[MASKED EMAIL]" in user_message @pytest.mark.asyncio