Addressing Heitor's feedback

Author: leandrodamascena

aws_lambda_powertools/utilities/serialization.py

@@ -4,6 +4,22 @@
 from typing import Any, Callable
 
 
+def base64_encode(data: str) -> str:
+    """Encode a string and returns Base64-encoded encoded value.
+
+    Parameters
+    ----------
+    data: str
+        The string to encode.
+
+    Returns
+    -------
+    str
+        The Base64-encoded encoded value.
+    """
+    return base64.b64encode(data.encode()).decode("utf-8")
+
+
 def base64_decode(data: str) -> str:
     """Decodes a Base64-encoded string and returns the decoded value.
 

tests/unit/data_classes/test_kinesis_firehose_response.py

@@ -4,6 +4,7 @@
     KinesisFirehoseDataTransformationResponse,
     KinesisFirehoseEvent,
 )
+from aws_lambda_powertools.utilities.serialization import base64_encode, base64_from_str
 from tests.functional.utils import load_event
 
 
@@ -42,30 +43,38 @@ def test_kinesis_firehose_response():
 
 
 def test_kinesis_firehose_create_response():
+    # GIVEN a Kinesis Firehose Event with two records
     raw_event = load_event("kinesisFirehoseKinesisEvent.json")
     parsed_event = KinesisFirehoseEvent(data=raw_event)
 
+    # WHEN we create a Data Transformation Response changing the data
+    # WHEN we add partitions keys
+
+    arbitrary_data = "arbitrary data"
+
     response = KinesisFirehoseDataTransformationResponse()
     for record in parsed_event.records:
-        # if data was delivered as json; caches loaded value
-        data = record.data_as_text
         metadata_partition = KinesisFirehoseDataTransformationRecordMetadata(partition_keys={"year": 2023})
         processed_record = record.build_data_transformation_response(
             result="Ok",
             metadata=metadata_partition,
+            data=base64_from_str(arbitrary_data),
         )
-        processed_record.data_from_text(data=data)
         response.add_record(record=processed_record)
-    response_dict = response.asdict()
-
-    res_records = list(response_dict["records"])
-    assert len(res_records) == 2
-    record_01, record_02 = res_records[:]
-    record01_raw = raw_event["records"][0]
-    assert record_01["result"] == "Ok"
-    assert record_01["recordId"] == record01_raw["recordId"]
-    assert record_01["data"] == record01_raw["data"]
-    assert record_01["metadata"]["partitionKeys"]["year"] == 2023
-
-    assert response.records[0].data_as_bytes == b"Hello World"
-    assert response.records[0].data_as_text == "Hello World"
+
+    # THEN we should have the same record data
+    record_01, record_02 = response.records[0], response.records[1]
+    raw_record_01, raw_record_02 = raw_event["records"][0], raw_event["records"][1]
+
+    assert len(response.records) == 2
+
+    assert record_01.result == "Ok"
+    assert record_02.result == "Ok"
+
+    assert record_01.record_id == raw_record_01["recordId"]
+    assert record_02.record_id == raw_record_02["recordId"]
+
+    assert record_01.data == base64_encode(arbitrary_data)
+    assert record_02.data == base64_encode(arbitrary_data)
+
+    assert record_01.metadata.partition_keys["year"] == 2023