[SB] Catch websocket exception (#34859)

* catch wb except * Update sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/_transport.py * eh
Azure · Apr 30, 2024 · 5fd2c52 · 5fd2c52
1 parent d2e868c
commit 5fd2c52
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_transport.py b/sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/_transport.py
@@ -771,7 +771,7 @@ def _read(self, n, initial=False, buffer=None, _errnos=None):  # pylint: disable
         :return: The data read.
         :rtype: bytearray
         """
-        from websocket import WebSocketTimeoutException
+        from websocket import WebSocketTimeoutException, WebSocketConnectionClosedException
         try:
             length = 0
             view = buffer or memoryview(bytearray(n))
@@ -794,6 +794,8 @@ def _read(self, n, initial=False, buffer=None, _errnos=None):  # pylint: disable
                 raise IOError("Websocket connection has already been closed.") from None
             except WebSocketTimeoutException as wte:
                 raise TimeoutError('Websocket receive timed out (%s)' % wte) from wte
+            except (WebSocketConnectionClosedException, SSLError) as e:
+                raise ConnectionError('Websocket disconnected: %r' % e) from e
         except:
             self._read_buffer = BytesIO(view[:length])
             raise

diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/_transport.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/_transport.py
@@ -771,7 +771,7 @@ def _read(self, n, initial=False, buffer=None, _errnos=None):  # pylint: disable
         :return: The data read.
         :rtype: bytearray
         """
-        from websocket import WebSocketTimeoutException
+        from websocket import WebSocketTimeoutException, WebSocketConnectionClosedException
         try:
             length = 0
             view = buffer or memoryview(bytearray(n))
@@ -794,6 +794,8 @@ def _read(self, n, initial=False, buffer=None, _errnos=None):  # pylint: disable
                 raise IOError("Websocket connection has already been closed.") from None
             except WebSocketTimeoutException as wte:
                 raise TimeoutError('Websocket receive timed out (%s)' % wte) from wte
+            except (WebSocketConnectionClosedException, SSLError) as e:
+                raise ConnectionError('Websocket disconnected: %r' % e) from e
         except:
             self._read_buffer = BytesIO(view[:length])
             raise