From 085b07cddeca59738fd9e3ffb173ec07f5a08b45 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 17 Dec 2025 14:04:33 +0530 Subject: [PATCH] Add test for handling gaurdrails param --- .../test_bedrock_image_gen_unit_tests.py | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/tests/image_gen_tests/test_bedrock_image_gen_unit_tests.py b/tests/image_gen_tests/test_bedrock_image_gen_unit_tests.py index 5526f22cd5e..73331547772 100644 --- a/tests/image_gen_tests/test_bedrock_image_gen_unit_tests.py +++ b/tests/image_gen_tests/test_bedrock_image_gen_unit_tests.py @@ -541,3 +541,28 @@ def test_amazon_titan_image_gen(): print(f"response cost: {response._hidden_params['response_cost']}") assert response._hidden_params["response_cost"] > 0 + + +def test_extract_headers_from_optional_params_with_guardrails(): + """Test that guardrail parameters are correctly extracted from optional_params and converted to headers""" + handler = BedrockImageGeneration() + + # Test with both guardrail parameters + optional_params = { + "guardrailIdentifier": "4cf5knqaeq15", + "guardrailVersion": "1", + "someOtherParam": "value", + } + + headers = handler._extract_headers_from_optional_params(optional_params) + + # Verify headers are correctly set + assert headers["x-amz-bedrock-guardrail-identifier"] == "4cf5knqaeq15" + assert headers["x-amz-bedrock-guardrail-version"] == "1" + + # Verify guardrail params are removed from optional_params + assert "guardrailIdentifier" not in optional_params + assert "guardrailVersion" not in optional_params + + # Verify other params remain in optional_params + assert optional_params["someOtherParam"] == "value"