Adding protobuf tests

Author: leandrodamascena

aws_lambda_powertools/utilities/kafka_consumer/deserializer/avro.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import io
-from typing import Any
 
 from avro.io import BinaryDecoder, DatumReader
 from avro.schema import parse as parse_schema
@@ -68,5 +67,5 @@ def deserialize(self, data: bytes | str) -> object:
             return self.reader.read(decoder)
         except Exception as e:
             raise KafkaConsumerDeserializationError(
-                f"Error trying to deserializer avro data - {type(e).__name__}: {str(e)}",
+                f"Error trying to deserialize avro data - {type(e).__name__}: {str(e)}",
             ) from e

aws_lambda_powertools/utilities/kafka_consumer/deserializer/base.py

@@ -2,7 +2,7 @@
 
 import base64
 from abc import ABC, abstractmethod
-from typing import Any, overload
+from typing import Any
 
 
 class DeserializerBase(ABC):
@@ -29,7 +29,7 @@ class DeserializerBase(ABC):
     """
 
     @abstractmethod
-    def deserialize(self, data: bytes | str) -> dict[str, Any] | str | object:
+    def deserialize(self, data: str) -> dict[str, Any] | str | object:
         """
         Deserialize input data to a Python dictionary.
 
@@ -38,8 +38,8 @@ def deserialize(self, data: bytes | str) -> dict[str, Any] | str | object:
 
         Parameters
         ----------
-        data : bytes or str
-            The data to deserialize, either as bytes or as a string.
+        data : str
+            The data to deserialize, it's always a base64 encoded string
 
         Returns
         -------
@@ -49,14 +49,9 @@ def deserialize(self, data: bytes | str) -> dict[str, Any] | str | object:
         raise NotImplementedError("Subclasses must implement the deserialize method")
 
     def _decode_input(self, data: bytes | str) -> bytes:
-        if isinstance(data, str):
+        try:
             return base64.b64decode(data)
-        elif isinstance(data, bytes):
-            return data
-        else:
-            try:
-                return base64.b64decode(data)
-            except Exception as e:
-                raise TypeError(
-                    f"Expected bytes or base64-encoded string, got {type(data).__name__}. Error: {str(e)}",
-                ) from e
+        except Exception as e:
+            raise TypeError(
+                f"Expected bytes or base64-encoded string, got {type(data).__name__}. Error: {str(e)}",
+            ) from e

aws_lambda_powertools/utilities/kafka_consumer/deserializer/default.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import base64
-from typing import Any
 
 from aws_lambda_powertools.utilities.kafka_consumer.deserializer.base import DeserializerBase
 

aws_lambda_powertools/utilities/kafka_consumer/deserializer/protobuf.py

@@ -61,5 +61,5 @@ def deserialize(self, data: bytes | str) -> dict:
             return MessageToDict(message, preserving_proto_field_name=True)
         except Exception as e:
             raise KafkaConsumerDeserializationError(
-                f"Protocol Buffer deserialization error: {type(e).__name__}: {str(e)}",
+                f"Error trying to deserialize protobuf data - {type(e).__name__}: {str(e)}",
             ) from e

tests/functional/kafka_consumer/_avro/test_kafka_consumer_with_avro.py

@@ -1,6 +1,7 @@
 import base64
 import io
 from copy import deepcopy
+from dataclasses import dataclass
 
 import pytest
 from avro.io import BinaryEncoder, DatumWriter
@@ -86,57 +87,58 @@ def kafka_event_with_avro_data(avro_encoded_value, avro_encoded_key):
     }
 
 
-def test_kafka_consumer_with_avro_and_dataclass(
-    kafka_event_with_avro_data,
-    avro_value_schema,
-    lambda_context,
-    user_value_dataclass,
-):
-    """Test Kafka consumer with Avro deserialization and dataclass output serialization."""
+@dataclass
+class UserValueDataClass:
+    name: str
+    age: int
+
+
+@dataclass
+class UserKeyClass:
+    user_id: str
+
+
+def test_kafka_consumer_with_avro(kafka_event_with_avro_data, avro_value_schema, lambda_context):
+    """Test Kafka consumer with Avro deserialization without output serialization."""
 
     # Create dict to capture results
     result_data = {}
 
-    schema_config = SchemaConfig(
-        value_schema_type="AVRO",
-        value_schema=avro_value_schema,
-        value_output_serializer=user_value_dataclass,
-    )
+    schema_config = SchemaConfig(value_schema_type="AVRO", value_schema=avro_value_schema)
 
     @kafka_consumer(schema_config=schema_config)
     def handler(event: ConsumerRecords, context):
         # Capture the results to verify
         record = next(event.records)
         result_data["value_type"] = type(record.value).__name__
-        result_data["name"] = record.value.name
-        result_data["age"] = record.value.age
+        result_data["name"] = record.value["name"]
+        result_data["age"] = record.value["age"]
         return {"processed": True}
 
     # Call the handler
     result = handler(kafka_event_with_avro_data, lambda_context)
 
     # Verify the results
     assert result == {"processed": True}
-    assert result_data["value_type"] == "UserValueDataClass"
+    assert result_data["value_type"] == "dict"
     assert result_data["name"] == "John Doe"
     assert result_data["age"] == 30
 
 
-def test_kafka_consumer_with_avro_and_custom_object(
+def test_kafka_consumer_with_avro_and_dataclass(
     kafka_event_with_avro_data,
     avro_value_schema,
     lambda_context,
-    user_value_dict,
 ):
-    """Test Kafka consumer with Avro deserialization and custom object serialization."""
+    """Test Kafka consumer with Avro deserialization and dataclass output serialization."""
 
     # Create dict to capture results
     result_data = {}
 
     schema_config = SchemaConfig(
         value_schema_type="AVRO",
         value_schema=avro_value_schema,
-        value_output_serializer=user_value_dict,
+        value_output_serializer=UserValueDataClass,
     )
 
     @kafka_consumer(schema_config=schema_config)
@@ -153,34 +155,43 @@ def handler(event: ConsumerRecords, context):
 
     # Verify the results
     assert result == {"processed": True}
-    assert result_data["value_type"] == "UserValueDict"
+    assert result_data["value_type"] == "UserValueDataClass"
     assert result_data["name"] == "John Doe"
     assert result_data["age"] == 30
 
 
-def test_kafka_consumer_with_avro_raw(kafka_event_with_avro_data, avro_value_schema, lambda_context):
-    """Test Kafka consumer with Avro deserialization without output serialization."""
+def test_kafka_consumer_with_avro_and_custom_object(
+    kafka_event_with_avro_data,
+    avro_value_schema,
+    lambda_context,
+    user_value_dict,
+):
+    """Test Kafka consumer with Avro deserialization and custom object serialization."""
 
     # Create dict to capture results
     result_data = {}
 
-    schema_config = SchemaConfig(value_schema_type="AVRO", value_schema=avro_value_schema)
+    schema_config = SchemaConfig(
+        value_schema_type="AVRO",
+        value_schema=avro_value_schema,
+        value_output_serializer=user_value_dict,
+    )
 
     @kafka_consumer(schema_config=schema_config)
     def handler(event: ConsumerRecords, context):
         # Capture the results to verify
         record = next(event.records)
         result_data["value_type"] = type(record.value).__name__
-        result_data["name"] = record.value["name"]
-        result_data["age"] = record.value["age"]
+        result_data["name"] = record.value.name
+        result_data["age"] = record.value.age
         return {"processed": True}
 
     # Call the handler
     result = handler(kafka_event_with_avro_data, lambda_context)
 
     # Verify the results
     assert result == {"processed": True}
-    assert result_data["value_type"] == "dict"
+    assert result_data["value_type"] == "UserValueDict"
     assert result_data["name"] == "John Doe"
     assert result_data["age"] == 30
 
@@ -207,7 +218,7 @@ def lambda_handler(event: ConsumerRecords, context):
 
     # The exact error message may vary depending on the Avro library's internals,
     # but should indicate a deserialization problem
-    assert "Error trying to deserializer avro data" in str(excinfo.value)
+    assert "Error trying to deserialize avro data" in str(excinfo.value)
 
 
 def test_kafka_consumer_with_invalid_avro_schema(kafka_event_with_avro_data, lambda_context):
@@ -245,8 +256,6 @@ def test_kafka_consumer_with_key_deserialization(
     lambda_context,
     avro_value_schema,
     avro_key_schema,
-    user_value_dataclass,
-    user_key_dataclass,
 ):
     """Test deserializing both key and value with different schemas and serializers."""
 
@@ -257,10 +266,10 @@ def test_kafka_consumer_with_key_deserialization(
     schema_config = SchemaConfig(
         value_schema_type="AVRO",
         value_schema=avro_value_schema,
-        value_output_serializer=user_value_dataclass,
+        value_output_serializer=UserValueDataClass,
         key_schema_type="AVRO",
         key_schema=avro_key_schema,
-        key_output_serializer=user_key_dataclass,
+        key_output_serializer=UserKeyClass,
     )
 
     @kafka_consumer(schema_config=schema_config)
@@ -283,48 +292,3 @@ def lambda_handler(event: ConsumerRecords, context):
     assert key_value_result["value_type"] == "UserValueDataClass"
     assert key_value_result["value_name"] == "John Doe"
     assert key_value_result["value_age"] == 30
-
-
-def test_kafka_consumer_with_different_serializers_for_key_and_value(
-    kafka_event_with_avro_data,
-    lambda_context,
-    avro_value_schema,
-    avro_key_schema,
-    user_key_dataclass,
-    user_value_dict,
-):
-    """Test using different serializer types for key and value."""
-
-    # Create dict to capture results
-    results = {}
-
-    # Create schema config with different serializers
-    schema_config = SchemaConfig(
-        value_schema_type="AVRO",
-        value_schema=avro_value_schema,
-        value_output_serializer=user_value_dict,
-        key_schema_type="AVRO",
-        key_schema=avro_key_schema,
-        key_output_serializer=user_key_dataclass,
-    )
-
-    @kafka_consumer(schema_config=schema_config)
-    def handler(event: ConsumerRecords, context):
-        record = next(event.records)
-        results["key_type"] = type(record.key).__name__
-        results["key_id"] = record.key.user_id
-        results["value_type"] = type(record.value).__name__
-        results["value_name"] = record.value.name
-        results["value_age"] = record.value.age
-        return {"processed": True}
-
-    # Call the handler
-    result = handler(kafka_event_with_avro_data, lambda_context)
-
-    # Verify the results
-    assert result == {"processed": True}
-    assert results["key_type"] == "UserKeyClass"
-    assert results["key_id"] == "user-123"
-    assert results["value_type"] == "UserValueDict"
-    assert results["value_name"] == "John Doe"
-    assert results["value_age"] == 30