Refactoring tests

Author: leandrodamascena

aws_lambda_powertools/utilities/data_classes/kafka_event.py

@@ -10,7 +10,7 @@
     from collections.abc import Iterator
 
 
-class ConsumerRecordSchemaMetadata(DictWrapper):
+class KafkaEventRecordSchemaMetadata(DictWrapper):
     @property
     def data_format(self) -> str | None:
         """The data format of the Kafka record."""
@@ -49,19 +49,19 @@ def timestamp_type(self) -> str:
         return self["timestampType"]
 
     @property
-    def key_schema_metadata(self) -> ConsumerRecordSchemaMetadata | None:
+    def key_schema_metadata(self) -> KafkaEventRecordSchemaMetadata | None:
         """The metadata of the Key Kafka record."""
         return (
-            None if self.get("keySchemaMetadata") is None else ConsumerRecordSchemaMetadata(self["keySchemaMetadata"])
+            None if self.get("keySchemaMetadata") is None else KafkaEventRecordSchemaMetadata(self["keySchemaMetadata"])
         )
 
     @property
-    def value_schema_metadata(self) -> ConsumerRecordSchemaMetadata | None:
+    def value_schema_metadata(self) -> KafkaEventRecordSchemaMetadata | None:
         """The metadata of the Value Kafka record."""
         return (
             None
             if self.get("valueSchemaMetadata") is None
-            else ConsumerRecordSchemaMetadata(self["valueSchemaMetadata"])
+            else KafkaEventRecordSchemaMetadata(self["valueSchemaMetadata"])
         )
 
 

aws_lambda_powertools/utilities/kafka/serialization/dataclass.py

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

aws_lambda_powertools/utilities/parser/models/kafka.py

@@ -8,6 +8,11 @@
 SERVERS_DELIMITER = ","
 
 
+class KafkaRecordSchemaMetadata(BaseModel):
+    dataFormat: str
+    schemaId: str
+
+
 class KafkaRecordModel(BaseModel):
     topic: str
     partition: int
@@ -17,6 +22,8 @@ class KafkaRecordModel(BaseModel):
     key: Optional[bytes] = None
     value: Union[str, Type[BaseModel]]
     headers: List[Dict[str, bytes]]
+    keySchemaMetadata: Optional[KafkaRecordSchemaMetadata] = None
+    valueSchemaMetadata: Optional[KafkaRecordSchemaMetadata] = None
 
     # key is optional; only decode if not None
     @field_validator("key", mode="before")

noxfile.py

@@ -67,6 +67,7 @@ def test_with_only_required_packages(session: nox.Session):
             f"{PREFIX_TESTS_UNIT}/data_classes/required_dependencies/",
             f"{PREFIX_TESTS_FUNCTIONAL}/event_handler/required_dependencies/",
             f"{PREFIX_TESTS_FUNCTIONAL}/batch/required_dependencies/",
+            f"{PREFIX_TESTS_FUNCTIONAL}/kafka_consumer/required_dependencies/",
         ],
     )
 
@@ -155,11 +156,13 @@ def test_with_pydantic_required_package(session: nox.Session):
     # Event Handler OpenAPI
     # Parser
     # Batch Processor with pydantic integration
+    # Kafka Consumer with Output to Pydantic
     build_and_run_test(
         session,
         folders=[
             f"{PREFIX_TESTS_FUNCTIONAL}/event_handler/_pydantic/",
             f"{PREFIX_TESTS_FUNCTIONAL}/batch/_pydantic/",
+            f"{PREFIX_TESTS_FUNCTIONAL}/kafka_consumer/_pydantic/",
             f"{PREFIX_TESTS_UNIT}/parser/_pydantic/",
             f"{PREFIX_TESTS_UNIT}/event_handler/_pydantic/",
         ],
@@ -196,3 +199,29 @@ def test_with_redis_and_boto3_sdk_as_required_package(session: nox.Session):
         ],
         extras="redis,aws-sdk",
     )
+
+
+@nox.session()
+def test_with_avro_required_package(session: nox.Session):
+    """Tests that only depends the Avro dependency"""
+    # Kafka Consumer with AVRO
+    build_and_run_test(
+        session,
+        folders=[
+            f"{PREFIX_TESTS_FUNCTIONAL}/kafka_consumer/_avro/",
+        ],
+        extras="kafka-consumer-avro",
+    )
+
+
+@nox.session()
+def test_with_protobuf_required_package(session: nox.Session):
+    """Tests that only depends the Protobuf dependency"""
+    # Kafka Consumer with PROTOBUF
+    build_and_run_test(
+        session,
+        folders=[
+            f"{PREFIX_TESTS_FUNCTIONAL}/kafka_consumer/_protobuf/",
+        ],
+        extras="kafka-consumer-protobuf",
+    )

tests/events/kafkaEventMsk.json

@@ -28,7 +28,15 @@
                     101
                  ]
               }
-           ]
+           ],
+           "valueSchemaMetadata": {
+              "dataFormat": "AVRO",
+              "schemaId": "1234"
+           },
+           "keySchemaMetadata": {
+              "dataFormat": "AVRO",
+              "schemaId": "1234"
+           }
         },
         {
            "topic":"mytopic",
@@ -53,7 +61,15 @@
                     101
                  ]
               }
-           ]
+           ],
+           "valueSchemaMetadata": {
+              "dataFormat": "AVRO",
+              "schemaId": "1234"
+           },
+           "keySchemaMetadata": {
+              "dataFormat": "AVRO",
+              "schemaId": "1234"
+           }
         },
         {
            "topic":"mytopic",
@@ -79,7 +95,15 @@
                     101
                  ]
               }
-           ]
+           ],
+           "valueSchemaMetadata": {
+              "dataFormat": "AVRO",
+              "schemaId": "1234"
+           },
+           "keySchemaMetadata": {
+              "dataFormat": "AVRO",
+              "schemaId": "1234"
+           }
         }
      ]
   }

tests/unit/data_classes/required_dependencies/test_kafka_event.py

@@ -34,6 +34,10 @@ def test_kafka_msk_event():
     assert record.json_value == {"key": "value"}
     assert record.decoded_headers == {"headerKey": b"headerValue"}
     assert record.decoded_headers["HeaderKey"] == b"headerValue"
+    assert record.key_schema_metadata.data_format == raw_record["keySchemaMetadata"]["dataFormat"]
+    assert record.key_schema_metadata.schema_id == raw_record["keySchemaMetadata"]["schemaId"]
+    assert record.value_schema_metadata.data_format == raw_record["valueSchemaMetadata"]["dataFormat"]
+    assert record.value_schema_metadata.schema_id == raw_record["valueSchemaMetadata"]["schemaId"]
 
     assert parsed_event.record == records[0]
     for i in range(1, 3):
@@ -68,6 +72,8 @@ def test_kafka_self_managed_event():
     assert record.json_value == {"key": "value"}
     assert record.decoded_headers == {"headerKey": b"headerValue"}
     assert record.decoded_headers["HeaderKey"] == b"headerValue"
+    assert record.key_schema_metadata is None
+    assert record.value_schema_metadata is None
 
     assert parsed_event.record == records[0]
 

tests/unit/parser/_pydantic/test_kafka.py

@@ -55,6 +55,8 @@ def test_self_managed_kafka_event():
     assert record.value == '{"key":"value"}'
     assert len(record.headers) == 1
     assert record.headers[0]["headerKey"] == b"headerValue"
+    assert record.keySchemaMetadata is None
+    assert record.valueSchemaMetadata is None
     record: KafkaRecordModel = records[1]
     assert record.key is None
 
@@ -82,6 +84,10 @@ def test_kafka_msk_event():
     assert record.value == '{"key":"value"}'
     assert len(record.headers) == 1
     assert record.headers[0]["headerKey"] == b"headerValue"
+    assert record.keySchemaMetadata.dataFormat == "AVRO"
+    assert record.keySchemaMetadata.schemaId == "1234"
+    assert record.valueSchemaMetadata.dataFormat == "AVRO"
+    assert record.valueSchemaMetadata.schemaId == "1234"
     for i in range(1, 3):
         record: KafkaRecordModel = records[i]
         assert record.key is None