Refactoring tests

Author: leandrodamascena

aws_lambda_powertools/utilities/kafka/consumer_records.py

@@ -50,7 +50,7 @@ def key(self) -> Any:
 
         if self.schema_config and self.schema_config.key_schema_type:
             schema_type = self.schema_config.key_schema_type
-            schema_str = self.schema_config.key_schema_str
+            schema_str = self.schema_config.key_schema
             output_serializer = self.schema_config.key_output_serializer
 
         # Always use get_deserializer if None it will default to DEFAULT
@@ -74,7 +74,7 @@ def value(self) -> Any:
 
         if self.schema_config and self.schema_config.value_schema_type:
             schema_type = self.schema_config.value_schema_type
-            schema_str = self.schema_config.value_schema_str
+            schema_str = self.schema_config.value_schema
             output_serializer = self.schema_config.value_output_serializer
 
         # Always use get_deserializer if None it will default to DEFAULT

aws_lambda_powertools/utilities/kafka/schema_config.py

@@ -77,8 +77,8 @@ def __init__(
             )
 
         self.value_schema_type = value_schema_type
-        self.value_schema_str = value_schema
+        self.value_schema = value_schema
         self.value_output_serializer = value_output_serializer
         self.key_schema_type = key_schema_type
-        self.key_schema_str = key_schema
+        self.key_schema = key_schema
         self.key_output_serializer = key_output_serializer

aws_lambda_powertools/utilities/kafka/serialization/base.py

@@ -4,6 +4,8 @@
 from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
     from aws_lambda_powertools.utilities.kafka.serialization.types import T
 
 
@@ -16,21 +18,21 @@ class OutputSerializerBase(ABC):
 
     Methods
     -------
-    serialize(data, output_class)
+    serialize(data, output)
         Abstract method that must be implemented by subclasses to serialize data.
 
     Examples
     --------
     >>> class MyOutputSerializer(OutputSerializerBase):
-    ...     def serialize(self, data: dict[str, Any], output_class=None):
-    ...         if output_class:
+    ...     def serialize(self, data: dict[str, Any], output=None):
+    ...         if output:
     ...             # Convert dictionary to class instance
-    ...             return output_class(**data)
+    ...             return output(**data)
     ...         return data  # Return as is if no output class provided
     """
 
     @abstractmethod
-    def serialize(self, data: dict[str, Any], output_class: type[T] | None = None) -> T | dict[str, Any]:
+    def serialize(self, data: dict[str, Any], output: type[T] | Callable | None = None) -> T | dict[str, Any]:
         """
         Serialize dictionary data into a specific output format or class instance.
 
@@ -41,14 +43,14 @@ def serialize(self, data: dict[str, Any], output_class: type[T] | None = None) -
         ----------
         data : dict[str, Any]
             The dictionary data to serialize.
-        output_class : type[T] or None, optional
+        output : type[T] or None, optional
             Optional class type to convert the dictionary into. If provided,
             the method should return an instance of this class.
 
         Returns
         -------
         T or dict[str, Any]
-            An instance of output_class if provided, otherwise a processed dictionary.
-            The generic type T represents the type of the output_class.
+            An instance of output if provided, otherwise a processed dictionary.
+            The generic type T represents the type of the output.
         """
         raise NotImplementedError("Subclasses must implement this method")

aws_lambda_powertools/utilities/kafka/serialization/custom_dict.py

@@ -2,23 +2,14 @@
 
 from typing import TYPE_CHECKING, Any
 
-from aws_lambda_powertools.utilities.kafka.exceptions import KafkaConsumerOutputSerializerError
 from aws_lambda_powertools.utilities.kafka.serialization.base import OutputSerializerBase
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
     from aws_lambda_powertools.utilities.kafka.serialization.types import T
 
 
 class CustomDictOutputSerializer(OutputSerializerBase):
-    def serialize(self, data: dict[str, Any], output_class: type[T] | None = None) -> T | dict[str, Any]:
-        if output_class is None:
-            return data
-
-        if not hasattr(output_class, "to_dict"):
-            raise KafkaConsumerOutputSerializerError("The output serialization class must have to_dict method")
-
-        # Instantiate and then populate
-        instance = output_class
-        for key, value in data.items():
-            setattr(instance, key, value)
-        return instance
+    def serialize(self, data: dict[str, Any], output: type[T] | Callable | None = None) -> T | dict[str, Any]:
+        return data if output is None else output(data)

aws_lambda_powertools/utilities/kafka/serialization/dataclass.py

@@ -1,18 +1,21 @@
 from __future__ import annotations
 
 from dataclasses import is_dataclass
-from typing import Any, cast
+from typing import TYPE_CHECKING, Any, cast
 
 from aws_lambda_powertools.utilities.kafka.serialization.base import OutputSerializerBase
 from aws_lambda_powertools.utilities.kafka.serialization.types import T
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
 
 class DataclassOutputSerializer(OutputSerializerBase):
-    def serialize(self, data: dict[str, Any], output_class: type[T] | None = None) -> T | dict[str, Any]:
-        if output_class is None:
+    def serialize(self, data: dict[str, Any], output: type[T] | Callable | None = None) -> T | dict[str, Any]:
+        if output is None:
             return data
 
-        if not is_dataclass(output_class):
+        if not is_dataclass(output):
             raise ValueError("Output class must be a dataclass")
 
-        return cast(T, output_class(**data))
+        return cast(T, output(**data))

aws_lambda_powertools/utilities/kafka/serialization/pydantic.py

@@ -7,14 +7,16 @@
 from aws_lambda_powertools.utilities.kafka.serialization.base import OutputSerializerBase
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
     from aws_lambda_powertools.utilities.kafka.serialization.types import T
 
 
 class PydanticOutputSerializer(OutputSerializerBase):
-    def serialize(self, data: dict[str, Any], output_class: type[T] | None = None) -> T | dict[str, Any]:
-        if output_class is None:
+    def serialize(self, data: dict[str, Any], output: type[T] | Callable | None = None) -> T | dict[str, Any]:
+        if output is None:
             return data
 
         # Use TypeAdapter for better support of Union types and other complex types
-        adapter = TypeAdapter(output_class)
+        adapter: TypeAdapter = TypeAdapter(output)
         return adapter.validate_python(data)

aws_lambda_powertools/utilities/kafka/serialization/serialization.py

@@ -7,23 +7,25 @@
 from aws_lambda_powertools.utilities.kafka.serialization.dataclass import DataclassOutputSerializer
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
     from aws_lambda_powertools.utilities.kafka.serialization.types import T
 
 
-def _get_output_serializer(output_class: type[T] | None = None) -> Any:
+def _get_output_serializer(output: type[T] | Callable | None = None) -> Any:
     """
     Returns the appropriate serializer for the given output class.
     Uses lazy imports to avoid unnecessary dependencies.
     """
-    if output_class is None:
+    if output is None:
         # Return a pass-through serializer if no output class is specified
         return CustomDictOutputSerializer()
 
     # Check if it's a dataclass
-    if is_dataclass(output_class):
+    if is_dataclass(output):
         return DataclassOutputSerializer()
 
-    if _is_pydantic_model(output_class):
+    if _is_pydantic_model(output):
         from aws_lambda_powertools.utilities.kafka.serialization.pydantic import PydanticOutputSerializer
 
         return PydanticOutputSerializer()
@@ -41,9 +43,12 @@ def _is_pydantic_model(obj: Any) -> bool:
     return False
 
 
-def serialize_to_output_type(data: object | dict[str, Any], output_class: type[T] | None = None) -> T | dict[str, Any]:
+def serialize_to_output_type(
+    data: object | dict[str, Any],
+    output: type[T] | Callable | None = None,
+) -> T | dict[str, Any]:
     """
     Helper function to directly serialize data to the specified output class
     """
-    serializer = _get_output_serializer(output_class)
-    return serializer.serialize(data, output_class)
+    serializer = _get_output_serializer(output)
+    return serializer.serialize(data, output)

tests/functional/kafka_consumer/_avro/test_kafka_consumer_with_avro.py

@@ -165,34 +165,34 @@ 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."""
 
+    def dict_output(data: dict) -> dict:
+        return data
+
     # 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=dict_output,
     )
 
     @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.get("name")
+        result_data["age"] = record.value.get("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"] == "UserValueDict"
     assert result_data["name"] == "John Doe"
     assert result_data["age"] == 30
 

tests/functional/kafka_consumer/_protobuf/test_kafka_consumer_with_protobuf.py

@@ -217,40 +217,33 @@ def test_kafka_consumer_with_custom_object(
     """Test Kafka consumer with Protobuf deserialization and custom object serialization."""
 
     # Define a custom output object class
-    class UserCustomObject:
-        def __init__(self, proto_message):
-            self.name = proto_message.name
-            self.age = proto_message.age
-            self.custom_field = f"{proto_message.name} is {proto_message.age} years old"
+    def dict_output(data: dict) -> dict:
+        return data
 
     # Create dict to capture results
     result_data = {}
 
     schema_config = SchemaConfig(
         value_schema_type="PROTOBUF",
         value_schema=User,
-        value_output_serializer=lambda msg: UserCustomObject(msg),
+        value_output_serializer=dict_output,
     )
 
     @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["custom_field"] = record.value.custom_field
+        result_data["name"] = record.value.get("name")
+        result_data["age"] = record.value.get("age")
         return {"processed": True}
 
     # Call the handler
     result = handler(kafka_event_with_proto_data, lambda_context)
 
     # Verify the results
     assert result == {"processed": True}
-    assert result_data["value_type"] == "UserCustomObject"
     assert result_data["name"] == "John Doe"
     assert result_data["age"] == 30
-    assert result_data["custom_field"] == "John Doe is 30 years old"
 
 
 def test_kafka_consumer_with_multiple_records(lambda_context):