Refactoring tests

Author: leandrodamascena

tests/functional/kafka_consumer/_pydantic/test_kafka_consumer_with_pydantic.py

@@ -54,38 +54,33 @@ class UserKeyModel(BaseModel):
 
 
 def test_kafka_consumer_with_json_value_and_pydantic(kafka_event_with_json_data, lambda_context):
-    """Test Kafka consumer with JSON deserialization and dataclass output serialization."""
-
-    # Create dict to capture results
-    result_data = {}
-
+    # GIVEN
+    # A Kafka consumer configured to deserialize JSON data
+    # and convert it to a Pydantic model instance
     schema_config = SchemaConfig(value_schema_type="JSON", value_output_serializer=UserValueModel)
 
     @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
-        return {"processed": True}
-
-    # Call the handler
+        # Extract the deserialized and serialized value
+        # which should be a UserValueModel instance
+        value: UserValueModel = event.record.value
+        return value
+
+    # WHEN
+    # The handler processes a Kafka event containing JSON-encoded data
+    # which is deserialized into a dictionary and then converted to a Pydantic model
     result = handler(kafka_event_with_json_data, lambda_context)
 
-    # Verify the results
-    assert result == {"processed": True}
-    assert result_data["value_type"] == "UserValueModel"
-    assert result_data["name"] == "John Doe"
-    assert result_data["age"] == 30
+    # THEN
+    # The result should be a UserValueModel instance with the correct properties
+    assert isinstance(result, UserValueModel)
+    assert result.name == "John Doe"
+    assert result.age == 30
 
 
 def test_kafka_consumer_with_json_value_and_union_tag(kafka_event_with_json_data, lambda_context):
     """Test Kafka consumer with JSON deserialization and dataclass output serialization."""
 
-    # Create dict to capture results
-    result_data = {}
-
     class UserValueModel(BaseModel):
         name: Literal["John Doe"]
         age: int
@@ -96,67 +91,73 @@ class UserValueModel2(BaseModel):
 
     UnionModel = Annotated[Union[UserValueModel, UserValueModel2], Field(discriminator="name")]
 
+    # GIVEN
+    # A Kafka consumer configured to deserialize JSON data
+    # and convert it to a Pydantic model instance with Union Tags
     schema_config = SchemaConfig(value_schema_type="JSON", value_output_serializer=UnionModel)
 
     @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
-        return {"processed": True}
-
-    # Call the handler
+        # Extract the deserialized and serialized value
+        # which should be a UserValueModel instance
+        value: UserValueModel = event.record.value
+        return value
+
+    # WHEN
+    # The handler processes a Kafka event containing JSON-encoded data
+    # which is deserialized into a dictionary and then converted to a Pydantic model
     result = handler(kafka_event_with_json_data, lambda_context)
 
-    # Verify the results
-    assert result == {"processed": True}
-    assert result_data["value_type"] == "UserValueModel"
-    assert result_data["name"] == "John Doe"
-    assert result_data["age"] == 30
+    # THEN
+    # The result should be a UserValueModel instance with the correct properties
+    assert isinstance(result, UserValueModel)
+    assert result.name == "John Doe"
+    assert result.age == 30
 
 
 def test_kafka_consumer_with_json_key_and_pydantic(kafka_event_with_json_data, lambda_context):
-    """Test Kafka consumer with JSON deserialization and dataclass output serialization."""
-
-    # Create dict to capture results
-    result_data = {}
-
-    schema_config = SchemaConfig(key_schema_type="JSON", key_output_serializer=UserKeyModel)
+    # GIVEN
+    # A Kafka consumer configured to deserialize only the key using JSON
+    # and convert it to a Pydantic UserKeyModel instance
+    schema_config = SchemaConfig(
+        key_schema_type="JSON",
+        key_output_serializer=UserKeyModel,
+    )
 
     @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.key).__name__
-        result_data["user_id"] = record.key.user_id
-        return {"processed": True}
+        # Extract the deserialized key to verify
+        key: UserKeyModel = event.record.key
+        return key
 
-    # Call the handler
+    # WHEN
+    # The handler processes a Kafka event, deserializing only the key portion as JSON
+    # while leaving the value in its original format
     result = handler(kafka_event_with_json_data, lambda_context)
 
-    # Verify the results
-    assert result == {"processed": True}
-    assert result_data["value_type"] == "UserKeyModel"
-    assert result_data["user_id"] == "123"
+    # THEN
+    # The key should be properly deserialized from JSON and converted to a UserKeyModel
+    # with the expected user_id value
+    assert isinstance(result, UserKeyModel)
+    assert result.user_id == "123"
 
 
-# Tests for Complex Types with Pydantic TypeAdapter
 def test_kafka_consumer_with_multiple_records(lambda_context):
-    """Test processing multiple records in a single event."""
-
-    # Create data for multiple records
+    # GIVEN
+    # Three different user records to process
+    # First user: John Doe, age 30
     data1 = {"name": "John Doe", "age": 30}
+    # Second user: Jane Smith, age 25
     data2 = {"name": "Jane Smith", "age": 25}
+    # Third user: Bob Johnson, age 40
     data3 = {"name": "Bob Johnson", "age": 40}
 
-    # Encode the data
+    # Base64-encoded JSON data for each record
     encoded1 = base64.b64encode(json.dumps(data1).encode("utf-8")).decode("utf-8")
     encoded2 = base64.b64encode(json.dumps(data2).encode("utf-8")).decode("utf-8")
     encoded3 = base64.b64encode(json.dumps(data3).encode("utf-8")).decode("utf-8")
 
-    # Create a kafka event with multiple records
+    # A Kafka event containing multiple records across different offsets
     multi_record_event = {
         "eventSource": "aws:kafka",
         "records": {
@@ -195,25 +196,31 @@ def test_kafka_consumer_with_multiple_records(lambda_context):
         },
     }
 
-    # Create schema config
-    schema_config = SchemaConfig(value_schema_type="JSON", value_output_serializer=UserValueModel)
-
-    # Create list to store processed records
+    # A list to capture processed record details
     processed_records = []
 
+    # A Kafka consumer configured to deserialize JSON and convert to Pydantic models
+    schema_config = SchemaConfig(value_schema_type="JSON", value_output_serializer=UserValueModel)
+
     @kafka_consumer(schema_config=schema_config)
     def handler(event: ConsumerRecords, context):
-        # Process all records
+        # Process each record and collect its properties
         for record in event.records:
             processed_records.append({"name": record.value.name, "age": record.value.age})
         return {"processed": len(processed_records)}
 
-    # Call the handler
+    # WHEN
+    # The handler processes the Kafka event containing multiple JSON records
     result = handler(multi_record_event, lambda_context)
 
-    # Verify the results
+    # THEN
+    # The handler should successfully process all three records
+    # and return the correct count
     assert result == {"processed": 3}
     assert len(processed_records) == 3
+
+    # All three users should be correctly deserialized and processed
+    # regardless of their order in the event
     assert any(r["name"] == "John Doe" and r["age"] == 30 for r in processed_records)
     assert any(r["name"] == "Jane Smith" and r["age"] == 25 for r in processed_records)
     assert any(r["name"] == "Bob Johnson" and r["age"] == 40 for r in processed_records)