Fix UnionModel handling on receiver side

Author: WiZeYAR

src/asyncapi_python/amqp/endpoint/receiver.py

@@ -22,6 +22,7 @@
 from .base import AbstractEndpoint, EndpointParams
 from ..error import Rejection, BadRequestRejection
 from ..operation import Operation
+from ..utils import UnionModel
 from aio_pika.abc import AbstractIncomingMessage, AbstractRobustQueue
 
 
@@ -80,6 +81,9 @@ async def _consumer(self, message: AbstractIncomingMessage):
     def _decode_payload(self, message: AbstractIncomingMessage) -> I:
         try:
             payload: I = self._params.decode(message.body, self._op.message_type)
+            # This trick is used to support union types in receivers
+            if isinstance(payload, UnionModel):
+                payload = cast(I, payload.root)
         except ValidationError as e:
             raise BadRequestRejection(e)
         return payload

src/asyncapi_python/amqp/utils.py

@@ -21,6 +21,9 @@
 U = TypeVar("U")
 
 
+class UnionModel(RootModel): ...
+
+
 def encode_message(message: T) -> bytes:
     return message.model_dump_json().encode()
 
@@ -30,6 +33,6 @@ def decode_message(message: bytes, schema: type[T]) -> T:
 
 
 @cache
-def union_model(types: tuple[type[U], ...]) -> type[RootModel[U]]:
+def union_model(types: tuple[type[U], ...]) -> type[UnionModel[U]]:
     UnionType = Union.__getitem__(types)
-    return RootModel[UnionType]  # type: ignore
+    return UnionModel[UnionType]  # type: ignore