add result warning, add asdict test, metadata test

Author: roger-zhangg

aws_lambda_powertools/utilities/data_classes/kinesis_firehose_event.py

@@ -1,5 +1,6 @@
 import base64
 import json
+import warnings
 from dataclasses import dataclass, field
 from typing import Any, Callable, Dict, Iterator, List, Optional
 
@@ -11,12 +12,24 @@
 @dataclass(repr=False, order=False, frozen=True)
 class KinesisFirehoseDataTransformationRecordMetadata:
     """
+    Metadata in Firehose Data Transform Record.
+
+    Parameters
+    ----------
+    partition_keys: Dict[str, str]
+        A dict of partition keys/value in string format, e.g. `{"year":"2023","month":"09"}`
+
     Documentation:
     --------------
     - https://docs.aws.amazon.com/firehose/latest/dev/dynamic-partitioning.html
     """
 
-    partition_keys: Dict[str, Any] = field(default_factory=lambda: {})
+    partition_keys: Dict[str, str] = field(default_factory=lambda: {})
+
+    def asdict(self) -> Dict:
+        if self.partition_keys is not None:
+            return {"partitionKeys": self.partition_keys}
+        return {}
 
 
 @dataclass(repr=False, order=False)
@@ -58,13 +71,19 @@ class KinesisFirehoseDataTransformationRecord:
     _json_data: Optional[Any] = None
 
     def asdict(self) -> Dict:
+        if self.result not in ["Ok", "Dropped", "ProcessingFailed"]:
+            warnings.warn(
+                stacklevel=1,
+                message=f'The result "{self.result}" is not valid, Choose from "Ok", "Dropped", "ProcessingFailed"',
+            )
+
         record: Dict[str, Any] = {
             "recordId": self.record_id,
             "result": self.result,
             "data": self.data,
         }
         if self.metadata:
-            record["metadata"] = self.metadata.__dict__
+            record["metadata"] = self.metadata.asdict()
         return record
 
     @property

tests/unit/data_classes/test_kinesis_firehose_response.py

@@ -8,6 +8,20 @@
 from tests.functional.utils import load_event
 
 
+def test_kinesis_firehose_response_metadata():
+    # When we create metadata with partition keys and attach to a firehose response record
+    metadata_partition = KinesisFirehoseDataTransformationRecordMetadata(partition_keys={"year": "2023"})
+
+    processed_record = KinesisFirehoseDataTransformationRecord(
+        record_id="test_id",
+        metadata=metadata_partition,
+        data="",
+    )
+    # Then we should have partition keys available in metadata field with same value
+    assert processed_record.metadata.partition_keys["year"] == "2023"
+    assert metadata_partition.asdict() == {"partitionKeys": {"year": "2023"}}
+
+
 def test_kinesis_firehose_response():
     # GIVEN a Kinesis Firehose Event with two records
     raw_event = load_event("kinesisFirehoseKinesisEvent.json")
@@ -16,10 +30,8 @@ def test_kinesis_firehose_response():
     # WHEN we create a Data Transformation Response without changing the data
     response = KinesisFirehoseDataTransformationResponse()
     for record in parsed_event.records:
-        metadata_partition = KinesisFirehoseDataTransformationRecordMetadata(partition_keys={"year": 2023})
         processed_record = KinesisFirehoseDataTransformationRecord(
             record_id=record.record_id,
-            metadata=metadata_partition,
             data=record.data,
         )
         response.add_record(record=processed_record)
@@ -39,7 +51,28 @@ def test_kinesis_firehose_response():
     assert record_01.data == raw_record_01["data"]
     assert record_02.data == raw_record_02["data"]
 
-    assert record_01.metadata.partition_keys["year"] == 2023
+
+def test_kinesis_firehose_response_asdict():
+    # 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"],
+    )
+    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,
+    )
+    response.add_record(record=processed_record)
+    assert response.asdict() == sample_response
 
 
 def test_kinesis_firehose_create_response():
@@ -54,7 +87,7 @@ def test_kinesis_firehose_create_response():
 
     response = KinesisFirehoseDataTransformationResponse()
     for record in parsed_event.records:
-        metadata_partition = KinesisFirehoseDataTransformationRecordMetadata(partition_keys={"year": 2023})
+        metadata_partition = KinesisFirehoseDataTransformationRecordMetadata(partition_keys={"year": "2023"})
         processed_record = record.build_data_transformation_response(
             result="Ok",
             metadata=metadata_partition,
@@ -77,4 +110,4 @@ def test_kinesis_firehose_create_response():
     assert record_01.data == base64_encode(arbitrary_data)
     assert record_02.data == base64_encode(arbitrary_data)
 
-    assert record_01.metadata.partition_keys["year"] == 2023
+    assert record_01.metadata.partition_keys["year"] == "2023"