refactor: addressing Heitor's feedback

Author: leandrodamascena

aws_lambda_powertools/utilities/data_classes/active_mq_event.py

@@ -1,4 +1,4 @@
-from typing import Any, Iterator, Optional
+from typing import Any, Dict, Iterator, Optional
 
 from aws_lambda_powertools.utilities.data_classes.common import DictWrapper
 from aws_lambda_powertools.utilities.data_classes.shared_functions import base64_decode
@@ -112,6 +112,10 @@ class ActiveMQEvent(DictWrapper):
     - https://aws.amazon.com/blogs/compute/using-amazon-mq-as-an-event-source-for-aws-lambda/
     """
 
+    def __init__(self, data: Dict[str, Any]):
+        super().__init__(data)
+        self._messages = (ActiveMQMessage(record) for record in self["messages"])
+
     @property
     def event_source(self) -> str:
         return self["eventSource"]
@@ -128,4 +132,4 @@ def messages(self) -> Iterator[ActiveMQMessage]:
 
     @property
     def message(self) -> ActiveMQMessage:
-        return next(self.messages)
+        return next(self._messages)

aws_lambda_powertools/utilities/data_classes/kafka_event.py

@@ -87,6 +87,11 @@ class KafkaEvent(DictWrapper):
     - https://docs.aws.amazon.com/lambda/latest/dg/with-msk.html
     """
 
+    def __init__(self, data: Dict[str, Any]):
+        super().__init__(data)
+        for chunk in self["records"].values():
+            self._records = (KafkaEventRecord(record) for record in chunk)
+
     @property
     def event_source(self) -> str:
         """The AWS service from which the Kafka event record originated."""
@@ -117,4 +122,4 @@ def records(self) -> Iterator[KafkaEventRecord]:
     @property
     def record(self) -> KafkaEventRecord:
         """The next Kafka record."""
-        return next(self.records)
+        return next(self._records)

tests/unit/data_classes/test_active_mq_event.py

@@ -1,3 +1,5 @@
+import pytest
+
 from aws_lambda_powertools.utilities.data_classes.active_mq_event import (
     ActiveMQEvent,
     ActiveMQMessage,
@@ -37,3 +39,20 @@ def test_active_mq_event():
     message = messages[1]
     assert message.json_data["timeout"] == 0
     assert message.json_data["data"] == "CZrmf0Gw8Ov4bqLQxD4E"
+
+    assert parsed_event.message == messages[1]
+    assert parsed_event.message == messages[2]
+
+
+def test_activemq_message_property_with_stopiteration_error():
+    # GIVEN a kafka event with one record
+    raw_event = load_event("activeMQEvent.json")
+    parsed_event = ActiveMQEvent(raw_event)
+
+    # WHEN calling record property four times
+    # THEN raise StopIteration
+    with pytest.raises(StopIteration):
+        assert parsed_event.message.message_id is not None
+        assert parsed_event.message.message_type is not None
+        assert parsed_event.message.data is not None
+        assert parsed_event.message.decoded_data is not None

tests/unit/data_classes/test_kafka_event.py

@@ -1,3 +1,5 @@
+import pytest
+
 from aws_lambda_powertools.utilities.data_classes import KafkaEvent
 from tests.functional.utils import load_event
 
@@ -31,6 +33,8 @@ def test_kafka_msk_event():
     assert record.decoded_headers == {"headerKey": b"headerValue"}
     assert record.get_header_value("HeaderKey", case_sensitive=False) == b"headerValue"
 
+    assert parsed_event.record == records[0]
+
 
 def test_kafka_self_managed_event():
     raw_event = load_event("kafkaEventSelfManaged.json")
@@ -59,3 +63,17 @@ def test_kafka_self_managed_event():
     assert record.json_value == {"key": "value"}
     assert record.decoded_headers == {"headerKey": b"headerValue"}
     assert record.get_header_value("HeaderKey", case_sensitive=False) == b"headerValue"
+
+    assert parsed_event.record == records[0]
+
+
+def test_kafka_record_property_with_stopiteration_error():
+    # GIVEN a kafka event with one record
+    raw_event = load_event("kafkaEventMsk.json")
+    parsed_event = KafkaEvent(raw_event)
+
+    # WHEN calling record property twice
+    # THEN raise StopIteration
+    with pytest.raises(StopIteration):
+        assert parsed_event.record.topic is not None
+        assert parsed_event.record.partition is not None