add more tests

Author: anafalcao

aws_lambda_powertools/event_handler/bedrock_agent.py

@@ -43,6 +43,7 @@ def __init__(
         self.prompt_session_attributes = prompt_session_attributes
         self.knowledge_bases_configuration = knowledge_bases_configuration
 
+
 class BedrockResponseBuilder(ResponseBuilder):
     """
     Bedrock Response Builder. This builds the response dict to be returned by Lambda when using Bedrock Agents.
@@ -52,6 +53,7 @@ class BedrockResponseBuilder(ResponseBuilder):
 
     @override
     def build(self, event: BedrockAgentEvent, *args) -> dict[str, Any]:
+        """Build the full response dict to be returned by the lambda"""
         self._route(event, None)
 
         bedrock_response = None

tests/functional/event_handler/_pydantic/test_bedrock_agent.py

@@ -240,3 +240,74 @@ def claims():
             },
         },
     ]
+
+
+def test_bedrock_agent_with_empty_bedrock_response():
+    # GIVEN a Bedrock Agent event
+    app = BedrockAgentResolver()
+
+    @app.get("/claims", description="Gets claims")
+    def claims():
+        return BedrockResponse(body={"message": "test"})
+
+    # WHEN calling the event handler
+    result = app(load_event("bedrockAgentEvent.json"), {})
+
+    # THEN process event correctly without optional attributes
+    assert result["messageVersion"] == "1.0"
+    assert result["response"]["httpStatusCode"] == 200
+    assert "sessionAttributes" not in result
+    assert "promptSessionAttributes" not in result
+    assert "knowledgeBasesConfiguration" not in result
+
+
+def test_bedrock_agent_with_partial_bedrock_response():
+    # GIVEN a Bedrock Agent event
+    app = BedrockAgentResolver()
+
+    @app.get("/claims", description="Gets claims")
+    def claims():
+        return BedrockResponse(
+            body={"message": "test"},
+            session_attributes={"user_id": "123"},
+            # Only include session_attributes to test partial response
+        )
+
+    # WHEN calling the event handler
+    result = app(load_event("bedrockAgentEvent.json"), {})
+
+    # THEN process event correctly with only session_attributes
+    assert result["messageVersion"] == "1.0"
+    assert result["response"]["httpStatusCode"] == 200
+    assert result["sessionAttributes"] == {"user_id": "123"}
+    assert "promptSessionAttributes" not in result
+    assert "knowledgeBasesConfiguration" not in result
+
+
+def test_bedrock_agent_with_different_attributes_combination():
+    # GIVEN a Bedrock Agent event
+    app = BedrockAgentResolver()
+
+    @app.get("/claims", description="Gets claims")
+    def claims():
+        return BedrockResponse(
+            body={"message": "test"},
+            prompt_session_attributes={"context": "testing"},
+            knowledge_bases_configuration=[
+                {
+                    "knowledgeBaseId": "kb-123",
+                    "retrievalConfiguration": {"vectorSearchConfiguration": {"numberOfResults": 3}},
+                },
+            ],
+            # Omit session_attributes to test different combination
+        )
+
+    # WHEN calling the event handler
+    result = app(load_event("bedrockAgentEvent.json"), {})
+
+    # THEN process event correctly with specific attributes
+    assert result["messageVersion"] == "1.0"
+    assert result["response"]["httpStatusCode"] == 200
+    assert "sessionAttributes" not in result
+    assert result["promptSessionAttributes"] == {"context": "testing"}
+    assert result["knowledgeBasesConfiguration"][0]["knowledgeBaseId"] == "kb-123"