Merge branch 'kinesis' of github.com:roger-zhangg/aws-lambda-powertools-python into kinesis

Author: roger-zhangg

aws_lambda_powertools/utilities/data_classes/kinesis_firehose_event.py

@@ -2,7 +2,7 @@
 import json
 import warnings
 from dataclasses import dataclass, field
-from typing import Any, Callable, Dict, Iterator, List, Optional
+from typing import Any, Callable, ClassVar, Dict, Iterator, List, Optional, Tuple
 
 from typing_extensions import Literal
 
@@ -62,6 +62,8 @@ class KinesisFirehoseDataTransformationRecord:
     - https://docs.aws.amazon.com/firehose/latest/dev/data-transformation.html
     """
 
+    _valid_result_types: ClassVar[Tuple[str, str, str]] = ("Ok", "Dropped", "ProcessingFailed")
+
     record_id: str
     result: Literal["Ok", "Dropped", "ProcessingFailed"] = "Ok"
     data: str = ""
@@ -71,7 +73,7 @@ class KinesisFirehoseDataTransformationRecord:
     _json_data: Optional[Any] = None
 
     def asdict(self) -> Dict:
-        if self.result not in ["Ok", "Dropped", "ProcessingFailed"]:
+        if self.result not in self._valid_result_types:
             warnings.warn(
                 stacklevel=1,
                 message=f'The result "{self.result}" is not valid, Choose from "Ok", "Dropped", "ProcessingFailed"',
@@ -132,7 +134,7 @@ def add_record(self, record: KinesisFirehoseDataTransformationRecord):
 
     def asdict(self) -> Dict:
         if not self.records:
-            raise ValueError("Kinesis Firehose doesn't accept empty response")
+            raise ValueError("Amazon Kinesis Data Firehose doesn't accept empty response")
 
         return {"records": [record.asdict() for record in self.records]}
 
@@ -216,7 +218,8 @@ def build_data_transformation_response(
         data: str = "",
         metadata: Optional[KinesisFirehoseDataTransformationRecordMetadata] = None,
     ) -> KinesisFirehoseDataTransformationRecord:
-        """create a KinesisFirehoseResponseRecord directly using the record_id and given values
+        """Create a KinesisFirehoseResponseRecord directly using the record_id and given values
+
         Parameters
         ----------
         result : Literal["Ok", "Dropped", "ProcessingFailed"]

tests/unit/data_classes/test_kinesis_firehose_response.py

@@ -53,24 +53,27 @@ def test_kinesis_firehose_response():
 
 
 def test_kinesis_firehose_response_asdict():
-    # Given the following example response provided by Firehose
+    # GIVEN the following example response provided by Firehose
     sample_response = {
         "records": [
             {"recordId": "sample_record", "data": "", "result": "Ok", "metadata": {"partitionKeys": {"year": "2023"}}},
         ],
     }
 
-    # Then asdict function should be able to return the same value
     response = KinesisFirehoseDataTransformationResponse()
     metadata_partition = KinesisFirehoseDataTransformationRecordMetadata(
         partition_keys=sample_response["records"][0]["metadata"]["partitionKeys"],
     )
+
+    # WHEN we create a transformation record with the exact same data
     processed_record = KinesisFirehoseDataTransformationRecord(
         record_id=sample_response["records"][0]["recordId"],
         data=sample_response["records"][0]["data"],
         result=sample_response["records"][0]["result"],
         metadata=metadata_partition,
     )
+
+    # THEN serialized response should return the same value
     response.add_record(record=processed_record)
     assert response.asdict() == sample_response