[SB] Catch websocket exception (#34859)

* catch wb except

* Update sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/_transport.py

* eh

Author: l0lawrence

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

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