From 2acedb48263166199f7c159443b955ccb31f6ca2 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 25 Oct 2025 12:59:11 -0700 Subject: [PATCH] test_bedrock_apply_guardrail_with_masking --- .../test_bedrock_apply_guardrail.py | 22 +++++++++---------- 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py b/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py index fbdf390faf..b30f0bf1e8 100644 --- a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py +++ b/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py @@ -1,17 +1,19 @@ """ Test the Bedrock guardrail apply_guardrail functionality """ -import sys import os -import pytest +import sys from unittest.mock import AsyncMock, Mock, patch +import pytest + sys.path.insert(0, os.path.abspath("../../../../..")) from fastapi import HTTPException -from litellm.types.guardrails import ApplyGuardrailRequest, ApplyGuardrailResponse + from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail +from litellm.types.guardrails import ApplyGuardrailRequest, ApplyGuardrailResponse @pytest.mark.asyncio @@ -95,11 +97,9 @@ async def test_bedrock_apply_guardrail_with_masking(): # Mock a response with masked content mock_response = { "action": "ALLOWED", - "content": [ + "outputs": [ { - "text": { - "text": "This is a test message with [REDACTED] content" - } + "text": "This is a test message with [REDACTED] content" } ] } @@ -145,7 +145,7 @@ async def test_bedrock_apply_guardrail_api_failure(): async def test_bedrock_apply_guardrail_endpoint_integration(): """Test the full endpoint integration with Bedrock guardrail""" from enterprise.litellm_enterprise.proxy.guardrails.endpoints import apply_guardrail - + # Create a real BedrockGuardrail instance guardrail = BedrockGuardrail( guardrail_name="test-bedrock-guard", @@ -160,11 +160,9 @@ async def test_bedrock_apply_guardrail_endpoint_integration(): # Mock a successful response from Bedrock mock_response = { "action": "ALLOWED", - "content": [ + "outputs": [ { - "text": { - "text": "This is a test message with processed content" - } + "text": "This is a test message with processed content" } ] }