[EventHubs] fix message type (Azure#28236)

* fix message type

* update EventData.message to hold uamqp Message for uamqp transport

* update dev reqs

* lint

Author: swathipil

sdk/eventhub/azure-eventhub/azure/eventhub/_common.py

@@ -54,17 +54,17 @@
     AmqpMessageProperties,
 )
 from ._pyamqp._message_backcompat import LegacyMessage, LegacyBatchMessage
-from ._pyamqp.message import Message
+from ._pyamqp.message import Message as pyamqp_Message
 from ._transport._pyamqp_transport import PyamqpTransport
 
 if TYPE_CHECKING:
     try:
         from uamqp import (  # pylint: disable=unused-import
-            Message as uamqp_Message,
+            Message,    # not importing as uamqp_Message, b/c type is exposed to user
             BatchMessage,
         )
     except ImportError:
-        uamqp_Message = None
+        Message = None
         BatchMessage = None
     from ._transport._base import AmqpTransport
 
@@ -134,8 +134,8 @@ def __init__(
         self._raw_amqp_message = AmqpAnnotatedMessage(  # type: ignore
             data_body=body, annotations={}, application_properties={}
         )
-        self._uamqp_message: Optional[Union[LegacyMessage, uamqp_Message]] = None
-        self._message: Message = None  # type: ignore
+        self._uamqp_message: Optional[Union[LegacyMessage, "Message"]] = None
+        self._message: Union["Message", pyamqp_Message] = None  # type: ignore
         self._raw_amqp_message.header = AmqpMessageHeader()
         self._raw_amqp_message.properties = AmqpMessageProperties()
         self.message_id = None
@@ -226,7 +226,7 @@ def from_message_content(  # pylint: disable=unused-argument
     @classmethod
     def _from_message(
         cls,
-        message: Union[uamqp_Message, Message],
+        message: Union["Message", pyamqp_Message],
         raw_amqp_message: Optional[AmqpAnnotatedMessage] = None,
     ) -> EventData:
         # pylint:disable=protected-access
@@ -260,7 +260,7 @@ def _decode_non_data_body_as_str(self, encoding: str = "UTF-8") -> str:
         return str(decode_with_recurse(seq_list, encoding))
 
     @property
-    def message(self) -> LegacyMessage:
+    def message(self) -> Union["Message", LegacyMessage]:
         """DEPRECATED: Get the underlying LegacyMessage.
          This is deprecated and will be removed in a later release.
 
@@ -278,7 +278,7 @@ def message(self) -> LegacyMessage:
         return self._uamqp_message
 
     @message.setter
-    def message(self, value: "uamqp_Message") -> None:
+    def message(self, value: "Message") -> None:
         """DEPRECATED: Set the underlying Message.
         This is deprecated and will be removed in a later release.
         """
@@ -632,7 +632,7 @@ def message(self) -> Union["BatchMessage", LegacyBatchMessage]:
             DeprecationWarning,
         )
         if not self._uamqp_message:
-            message = AmqpAnnotatedMessage(message=Message(*self._message))
+            message = AmqpAnnotatedMessage(message=pyamqp_Message(*self._message))
             self._uamqp_message = LegacyBatchMessage(
                 message,
                 to_outgoing_amqp_message=PyamqpTransport().to_outgoing_amqp_message,

sdk/eventhub/azure-eventhub/azure/eventhub/_consumer.py

@@ -13,6 +13,7 @@
 from ._common import EventData
 from ._client_base import ConsumerProducerMixin
 from ._utils import create_properties, event_position_selector
+from ._transport._pyamqp_transport import PyamqpTransport
 from ._constants import (
     EPOCH_SYMBOL,
     TIMEOUT_SYMBOL,
@@ -182,6 +183,8 @@ def _next_message_in_buffer(self):
         # pylint:disable=protected-access
         message = self._message_buffer.popleft()
         event_data = EventData._from_message(message)
+        if self._amqp_transport != PyamqpTransport:
+            event_data._uamqp_message == message    # pylint: disable=pointless-statement
         self._last_received_event = event_data
         return event_data
 

sdk/eventhub/azure-eventhub/azure/eventhub/aio/_consumer_async.py

@@ -177,10 +177,12 @@ def _create_handler(
     async def _open_with_retry(self) -> None:
         await self._do_retryable_operation(self._open, operation_need_param=False)
 
+    # only used by _uamqp_transport_async
     def _next_message_in_buffer(self):
         # pylint:disable=protected-access
         message = self._message_buffer.popleft()
         event_data = EventData._from_message(message)
+        event_data._uamqp_message = message
         self._last_received_event = event_data
         return event_data
 

sdk/eventhub/azure-eventhub/dev_requirements.txt

@@ -3,7 +3,7 @@
 -e ../../identity/azure-identity
 -e ../azure-mgmt-eventhub
 azure-mgmt-resource==20.0.0
-aiohttp>=3.0
-websocket-client
+aiohttp>=3.0,<4.0
+websocket-client==1.4.2
 -e ../../../tools/azure-devtools
 uamqp>=1.6.3,<2.0.0

sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_receive_async.py

@@ -8,9 +8,15 @@
 import pytest
 import time
 
+try:
+    import uamqp
+except (ModuleNotFoundError, ImportError):
+    uamqp = None
+
 from azure.eventhub import EventData, TransportType
 from azure.eventhub.exceptions import EventHubError
 from azure.eventhub.aio import EventHubProducerClient, EventHubConsumerClient
+from azure.eventhub._pyamqp._message_backcompat import LegacyMessage
 
 
 @pytest.mark.liveTest
@@ -27,6 +33,10 @@ async def on_event(partition_context, event):
             assert ", sequence_number: " in event_str
             assert ", enqueued_time: " in event_str
             assert ", partition_key: 0" in event_str
+        if uamqp_transport:
+            assert isinstance(event.message, uamqp.Message)
+        else:
+            assert isinstance(event.message, LegacyMessage)
 
     on_event.called = False
     connection_str, senders = connstr_senders

sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_receive.py

@@ -8,8 +8,14 @@
 import pytest
 import time
 
+try:
+    import uamqp
+except (ModuleNotFoundError, ImportError):
+    uamqp = None
+
 from azure.eventhub import EventData, TransportType, EventHubConsumerClient
 from azure.eventhub.exceptions import EventHubError
+from azure.eventhub._pyamqp._message_backcompat import LegacyMessage
 
 
 @pytest.mark.liveTest
@@ -25,6 +31,11 @@ def on_event(partition_context, event):
             assert ", sequence_number: " in event_str
             assert ", enqueued_time: " in event_str
             assert ", partition_key: 0" in event_str
+        if uamqp_transport:
+            assert isinstance(event.message, uamqp.Message)
+        else:
+            assert isinstance(event.message, LegacyMessage)
+
     on_event.called = False
     connection_str, senders = connstr_senders
     client = EventHubConsumerClient.from_connection_string(