[Azure AMQP] Handle Websocket Disconnect in Async (Azure#32061)

* handle websocket close

* make sure buffer has left over bytes

Author: kashifkhan

sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/aio/_transport_async.py

@@ -499,18 +499,23 @@ async def _read(self, toread, buffer=None, **kwargs):  # pylint: disable=unused-
         length += nbytes
         toread -= nbytes
         try:
-            while toread:
-                data = await self.sock.receive_bytes()
-                read_length = len(data)
-                if read_length <= toread:
-                    view[length : length + read_length] = data
-                    toread -= read_length
-                    length += read_length
-                else:
-                    view[length : length + toread] = data[0:toread]
-                    self._read_buffer = BytesIO(data[toread:])
-                    toread = 0
-            return view
+            try:
+                while toread:
+                    data = await self.sock.receive_bytes()
+                    read_length = len(data)
+                    if read_length <= toread:
+                        view[length : length + read_length] = data
+                        toread -= read_length
+                        length += read_length
+                    else:
+                        view[length : length + toread] = data[0:toread]
+                        self._read_buffer = BytesIO(data[toread:])
+                        toread = 0
+                return view
+            except TypeError as te:
+                # aiohttp websocket raises TypeError when a websocket disconnects, as it ends up
+                # reading None over the wire and cant convert to bytes.
+                raise ConnectionError("Websocket disconnected: %r" % te) from None
         except:
             self._read_buffer = BytesIO(view[:length])
             raise

sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/aio/_transport_async.py

@@ -499,18 +499,23 @@ async def _read(self, toread, buffer=None, **kwargs):  # pylint: disable=unused-
         length += nbytes
         toread -= nbytes
         try:
-            while toread:
-                data = await self.sock.receive_bytes()
-                read_length = len(data)
-                if read_length <= toread:
-                    view[length : length + read_length] = data
-                    toread -= read_length
-                    length += read_length
-                else:
-                    view[length : length + toread] = data[0:toread]
-                    self._read_buffer = BytesIO(data[toread:])
-                    toread = 0
-            return view
+            try:
+                while toread:
+                    data = await self.sock.receive_bytes()
+                    read_length = len(data)
+                    if read_length <= toread:
+                        view[length : length + read_length] = data
+                        toread -= read_length
+                        length += read_length
+                    else:
+                        view[length : length + toread] = data[0:toread]
+                        self._read_buffer = BytesIO(data[toread:])
+                        toread = 0
+                return view
+            except TypeError as te:
+                # aiohttp websocket raises TypeError when a websocket disconnects, as it ends up
+                # reading None over the wire and cant convert to bytes.
+                raise ConnectionError("Websocket disconnected: %r" % te) from None
         except:
             self._read_buffer = BytesIO(view[:length])
             raise