fix: add tests

Author: rubenfonseca

Diff for: aws_lambda_powertools/utilities/data_classes/bedrock_agent_event.py

@@ -47,7 +47,35 @@ def content(self) -> Dict[str, BedrockAgentRequestMedia]:
         return {k: BedrockAgentRequestMedia(v) for k, v in self["content"].items()}
 
 
+class BedrockAgentActionGroup(DictWrapper):
+    @property
+    def action_group(self) -> str:
+        return self["actionGroup"]
+
+    @property
+    def api_path(self) -> str:
+        return self["apiPath"]
+
+    @property
+    def http_method(self) -> str:
+        return self["httpMethod"]
+
+    @property
+    def parameters(self) -> List[BedrockAgentProperty]:
+        return [BedrockAgentProperty(x) for x in self["parameters"]]
+
+    @property
+    def request_body(self) -> BedrockAgentRequestBody:
+        return BedrockAgentRequestBody(self["requestBody"])
+
+
 class BedrockAgentEvent(DictWrapper):
+    """
+    Bedrock Agent input event
+
+    See https://docs.aws.amazon.com/bedrock/latest/userguide/agents-create.html
+    """
+
     @property
     def message_version(self) -> str:
         return self["messageVersion"]
@@ -61,20 +89,12 @@ def session_id(self) -> str:
         return self["sessionId"]
 
     @property
-    def action_group(self) -> str:
-        return self["actionGroup"]
+    def action_groups(self) -> List[BedrockAgentActionGroup]:
+        return [BedrockAgentActionGroup(x) for x in self["actionGroups"]]
 
     @property
-    def api_path(self) -> str:
-        return self["apiPath"]
-
-    @property
-    def http_method(self) -> str:
-        return self["httpMethod"]
-
-    @property
-    def parameters(self) -> List[BedrockAgentProperty]:
-        return [BedrockAgentProperty(x) for x in self["parameters"]]
+    def agent(self) -> BedrockAgentInfo:
+        return BedrockAgentInfo(self["agent"])
 
     @property
     def session_attributes(self) -> Dict[str, str]:
@@ -110,10 +130,16 @@ def response_body(self) -> Dict[str, BedrockAgentResponseMedia]:
 
 
 class BedrockAgentResponseEvent(DictWrapper):
+    """
+    Bedrock Agent output event
+
+    See: https://docs.aws.amazon.com/bedrock/latest/userguide/agents-create.html
+    """
+
     @property
     def message_version(self) -> str:
         return self["messageVersion"]
 
     @property
-    def response(self) -> BedrockAgentResponsePayload:
-        return BedrockAgentResponsePayload(self["response"])
+    def responses(self) -> List[BedrockAgentResponsePayload]:
+        return [BedrockAgentResponsePayload(x) for x in self["response"]]

Diff for: tests/events/bedrockAgentEvent.json

@@ -0,0 +1,26 @@
+{
+  "sessionId": "12345678912345",
+  "agent": {
+    "alias": "TSTALIASID",
+    "name": "test",
+    "version": "DRAFT",
+    "id": "8ZXY0W8P1H"
+  },
+  "inputText": "ABC12345",
+  "sessionAttributes": {},
+  "messageVersion": "1.0",
+  "actionGroups": [
+    {
+      "httpMethod": "GET",
+      "apiPath": "/claims/{claimId}/identify-missing-documents",
+      "actionGroup": "ClaimManagementActionGroup",
+      "parameters": [
+        {
+          "name": "claimId",
+          "type": "string",
+          "value": "ABC12345"
+        }
+      ]
+    }
+  ]
+}

Diff for: tests/events/bedrockAgentResponseEvent.json

@@ -0,0 +1,16 @@
+{
+  "messageVersion": "1.0",
+  "response": [
+    {
+      "actionGroup": "ClaimManagementActionGroup",
+      "apiPath": "/claims/{claimId}/identify-missing-documents",
+      "httpMethod": "GET",
+      "httpStatusCode": 200,
+      "responseBody": {
+        "application/json": {
+          "body": "This is the response"
+        }
+      }
+    }
+  ]
+}

Diff for: tests/unit/data_classes/test_bedrock_agent_event.py

@@ -0,0 +1,56 @@
+from aws_lambda_powertools.utilities.data_classes import BedrockAgentEvent, BedrockAgentResponseEvent
+from tests.functional.utils import load_event
+
+
+def test_bedrock_agent_event():
+    raw_event = load_event("bedrockAgentEvent.json")
+    parsed_event = BedrockAgentEvent(raw_event)
+
+    assert parsed_event.session_id == "12345678912345"
+    assert parsed_event.input_text == "ABC12345"
+    assert parsed_event.session_attributes == {}
+    assert parsed_event.message_version == "1.0"
+
+    agent = parsed_event.agent
+    assert agent.alias == "TSTALIASID"
+    assert agent.name == "test"
+    assert agent.version == "DRAFT"
+    assert agent.id == "8ZXY0W8P1H"
+
+    action_groups = parsed_event.action_groups
+    assert len(action_groups) == 1
+
+    action_group = action_groups[0]
+    assert action_group.http_method == "GET"
+    assert action_group.api_path == "/claims/{claimId}/identify-missing-documents"
+    assert action_group.action_group == "ClaimManagementActionGroup"
+
+    parameters = action_group.parameters
+    assert len(parameters) == 1
+
+    parameter = parameters[0]
+    assert parameter.name == "claimId"
+    assert parameter.type == "string"
+    assert parameter.value == "ABC12345"
+
+
+def test_bedrock_agent_response_event():
+    raw_event = load_event("bedrockAgentResponseEvent.json")
+    parsed_event = BedrockAgentResponseEvent(raw_event)
+
+    assert parsed_event.message_version == "1.0"
+
+    responses = parsed_event.responses
+    assert len(responses) == 1
+
+    response = responses[0]
+    assert response.action_group == "ClaimManagementActionGroup"
+    assert response.api_path == "/claims/{claimId}/identify-missing-documents"
+    assert response.http_method == "GET"
+    assert response.http_status_code == 200
+
+    response_body = response.response_body
+    assert "application/json" in response_body
+
+    json_response = response_body["application/json"]
+    assert json_response.body == "This is the response"