Skip to content

Commit 917c8bb

Browse files
ARajan1084ARajan1084
committed
Update test_bedrock_guardrails.py
1 parent a86b9a1 commit 917c8bb

File tree

1 file changed

+191
-0
lines changed

1 file changed

+191
-0
lines changed

‎tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py‎

Lines changed: 191 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -239,6 +239,7 @@ async def test_bedrock_guardrail_logging_uses_redacted_response():
239239
) as mock_load_creds, patch.object(
240240
guardrail, "_prepare_request", return_value=MagicMock()
241241
) as mock_prepare_request:
242+
242243
mock_post.return_value = mock_bedrock_response
243244

244245
# Call the method that should log the redacted response
@@ -345,6 +346,7 @@ async def test_bedrock_guardrail_original_response_not_modified():
345346
) as mock_load_creds, patch.object(
346347
guardrail, "_prepare_request", return_value=MagicMock()
347348
) as mock_prepare_request:
349+
348350
mock_post.return_value = mock_bedrock_response
349351

350352
# Call the method
@@ -1049,6 +1051,195 @@ async def test_bedrock_guardrail_parameter_takes_precedence_over_env(monkeypatch
10491051
print(f"Parameter precedence test passed. URL: {prepped_request.url}")
10501052

10511053

1054+
@pytest.mark.asyncio
1055+
async def test_bedrock_guardrail_respects_custom_runtime_endpoint(monkeypatch):
1056+
"""Test that BedrockGuardrail respects aws_bedrock_runtime_endpoint when set"""
1057+
1058+
# Clear any existing environment variable to ensure clean test
1059+
monkeypatch.delenv("AWS_BEDROCK_RUNTIME_ENDPOINT", raising=False)
1060+
1061+
# Create guardrail with custom runtime endpoint
1062+
custom_endpoint = "https://custom-bedrock.example.com"
1063+
guardrail = BedrockGuardrail(
1064+
guardrailIdentifier="test-guardrail",
1065+
guardrailVersion="DRAFT",
1066+
aws_bedrock_runtime_endpoint=custom_endpoint,
1067+
)
1068+
1069+
# Mock credentials
1070+
mock_credentials = MagicMock()
1071+
mock_credentials.access_key = "test-access-key"
1072+
mock_credentials.secret_key = "test-secret-key"
1073+
mock_credentials.token = None
1074+
1075+
# Test data
1076+
data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]}
1077+
optional_params = {}
1078+
aws_region_name = "us-east-1"
1079+
1080+
# Mock the _load_credentials method to avoid actual AWS credential loading
1081+
with patch.object(
1082+
guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name)
1083+
):
1084+
# Call _prepare_request which internally calls get_runtime_endpoint
1085+
prepped_request = guardrail._prepare_request(
1086+
credentials=mock_credentials,
1087+
data=data,
1088+
optional_params=optional_params,
1089+
aws_region_name=aws_region_name,
1090+
)
1091+
1092+
# Verify that the custom endpoint is used in the URL
1093+
expected_url = f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply"
1094+
assert (
1095+
prepped_request.url == expected_url
1096+
), f"Expected URL to contain custom endpoint. Got: {prepped_request.url}"
1097+
1098+
print(f"Custom runtime endpoint test passed. URL: {prepped_request.url}")
1099+
1100+
1101+
@pytest.mark.asyncio
1102+
async def test_bedrock_guardrail_respects_env_runtime_endpoint(monkeypatch):
1103+
"""Test that BedrockGuardrail respects AWS_BEDROCK_RUNTIME_ENDPOINT environment variable"""
1104+
1105+
custom_endpoint = "https://env-bedrock.example.com"
1106+
1107+
# Set the environment variable
1108+
monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", custom_endpoint)
1109+
1110+
# Create guardrail without explicit aws_bedrock_runtime_endpoint
1111+
guardrail = BedrockGuardrail(
1112+
guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT"
1113+
)
1114+
1115+
# Mock credentials
1116+
mock_credentials = MagicMock()
1117+
mock_credentials.access_key = "test-access-key"
1118+
mock_credentials.secret_key = "test-secret-key"
1119+
mock_credentials.token = None
1120+
1121+
# Test data
1122+
data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]}
1123+
optional_params = {}
1124+
aws_region_name = "us-east-1"
1125+
1126+
# Mock the _load_credentials method
1127+
with patch.object(
1128+
guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name)
1129+
):
1130+
# Call _prepare_request which internally calls get_runtime_endpoint
1131+
prepped_request = guardrail._prepare_request(
1132+
credentials=mock_credentials,
1133+
data=data,
1134+
optional_params=optional_params,
1135+
aws_region_name=aws_region_name,
1136+
)
1137+
1138+
# Verify that the custom endpoint from environment is used in the URL
1139+
expected_url = f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply"
1140+
assert (
1141+
prepped_request.url == expected_url
1142+
), f"Expected URL to contain env endpoint. Got: {prepped_request.url}"
1143+
1144+
print(f"Environment runtime endpoint test passed. URL: {prepped_request.url}")
1145+
1146+
1147+
@pytest.mark.asyncio
1148+
async def test_bedrock_guardrail_uses_default_endpoint_when_no_custom_set(monkeypatch):
1149+
"""Test that BedrockGuardrail uses default endpoint when no custom endpoint is set"""
1150+
1151+
# Ensure no environment variable is set
1152+
monkeypatch.delenv("AWS_BEDROCK_RUNTIME_ENDPOINT", raising=False)
1153+
1154+
# Create guardrail without any custom endpoint
1155+
guardrail = BedrockGuardrail(
1156+
guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT"
1157+
)
1158+
1159+
# Mock credentials
1160+
mock_credentials = MagicMock()
1161+
mock_credentials.access_key = "test-access-key"
1162+
mock_credentials.secret_key = "test-secret-key"
1163+
mock_credentials.token = None
1164+
1165+
# Test data
1166+
data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]}
1167+
optional_params = {}
1168+
aws_region_name = "us-west-2"
1169+
1170+
# Mock the _load_credentials method
1171+
with patch.object(
1172+
guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name)
1173+
):
1174+
# Call _prepare_request which internally calls get_runtime_endpoint
1175+
prepped_request = guardrail._prepare_request(
1176+
credentials=mock_credentials,
1177+
data=data,
1178+
optional_params=optional_params,
1179+
aws_region_name=aws_region_name,
1180+
)
1181+
1182+
# Verify that the default endpoint is used
1183+
expected_url = f"https://bedrock-runtime.{aws_region_name}.amazonaws.com/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply"
1184+
assert (
1185+
prepped_request.url == expected_url
1186+
), f"Expected default URL. Got: {prepped_request.url}"
1187+
1188+
print(f"Default endpoint test passed. URL: {prepped_request.url}")
1189+
1190+
1191+
@pytest.mark.asyncio
1192+
async def test_bedrock_guardrail_parameter_takes_precedence_over_env(monkeypatch):
1193+
"""Test that aws_bedrock_runtime_endpoint parameter takes precedence over environment variable
1194+
1195+
This test verifies the corrected behavior where the parameter should take precedence
1196+
over the environment variable, consistent with the endpoint_url logic.
1197+
"""
1198+
1199+
param_endpoint = "https://param-bedrock.example.com"
1200+
env_endpoint = "https://env-bedrock.example.com"
1201+
1202+
# Set environment variable
1203+
monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", env_endpoint)
1204+
1205+
# Create guardrail with explicit aws_bedrock_runtime_endpoint
1206+
guardrail = BedrockGuardrail(
1207+
guardrailIdentifier="test-guardrail",
1208+
guardrailVersion="DRAFT",
1209+
aws_bedrock_runtime_endpoint=param_endpoint,
1210+
)
1211+
1212+
# Mock credentials
1213+
mock_credentials = MagicMock()
1214+
mock_credentials.access_key = "test-access-key"
1215+
mock_credentials.secret_key = "test-secret-key"
1216+
mock_credentials.token = None
1217+
1218+
# Test data
1219+
data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]}
1220+
optional_params = {}
1221+
aws_region_name = "us-east-1"
1222+
1223+
# Mock the _load_credentials method
1224+
with patch.object(
1225+
guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name)
1226+
):
1227+
# Call _prepare_request which internally calls get_runtime_endpoint
1228+
prepped_request = guardrail._prepare_request(
1229+
credentials=mock_credentials,
1230+
data=data,
1231+
optional_params=optional_params,
1232+
aws_region_name=aws_region_name,
1233+
)
1234+
1235+
# Verify that the parameter takes precedence over environment variable
1236+
expected_url = f"{param_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply"
1237+
assert (
1238+
prepped_request.url == expected_url
1239+
), f"Expected parameter endpoint to take precedence. Got: {prepped_request.url}"
1240+
1241+
print(f"Parameter precedence test passed. URL: {prepped_request.url}")
1242+
10521243
@pytest.mark.asyncio
10531244
async def test_bedrock_guardrail_200_with_exception_in_output_raises_and_logs_failure():
10541245
"""

0 commit comments

Comments
 (0)