Refactoring tests

Author: leandrodamascena

aws_lambda_powertools/utilities/kafka/serialization/serialization.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
+import sys
 from dataclasses import is_dataclass
-from typing import TYPE_CHECKING, Any
+from types import UnionType
+from typing import TYPE_CHECKING, Annotated, Any, Optional, Union, get_args, get_origin
 
 from aws_lambda_powertools.utilities.kafka.serialization.custom_dict import CustomDictOutputSerializer
 from aws_lambda_powertools.utilities.kafka.serialization.dataclass import DataclassOutputSerializer
@@ -17,10 +19,6 @@ 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 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):
         return DataclassOutputSerializer()
@@ -40,6 +38,14 @@ def _is_pydantic_model(obj: Any) -> bool:
         has_model_fields = getattr(obj, "model_fields", None) is not None
         has_model_validate = callable(getattr(obj, "model_validate", None))
         return has_model_fields and has_model_validate
+
+    origin = get_origin(obj)
+    if origin in (Union, Optional, Annotated) or (sys.version_info >= (3, 10) and origin in (Union, UnionType)):
+        # Check if any element in the Union is a Pydantic model
+        for arg in get_args(obj):
+            if _is_pydantic_model(arg):
+                return True
+
     return False
 
 

tests/functional/kafka_consumer/_pydantic/test_kafka_consumer_with_pydantic.py

@@ -1,6 +1,6 @@
 import base64
 import json
-from typing import Literal, Union
+from typing import Annotated, Literal, Union
 
 import pytest
 from pydantic import BaseModel, Field
@@ -94,10 +94,9 @@ class UserValueModel2(BaseModel):
         name: Literal["Not using"]
         email: str
 
-    class Model(BaseModel):
-        name: Union[UserValueModel, UserValueModel2] = Field(discriminator="name")
+    UnionModel = Annotated[Union[UserValueModel, UserValueModel2], Field(discriminator="name")]
 
-    schema_config = SchemaConfig(value_schema_type="JSON", value_output_serializer=UserValueModel)
+    schema_config = SchemaConfig(value_schema_type="JSON", value_output_serializer=UnionModel)
 
     @kafka_consumer(schema_config=schema_config)
     def handler(event: ConsumerRecords, context):