clean up

Author: dmulcahey

zha/application/gateway.py

@@ -911,7 +911,7 @@ def __init__(self, config: ZHAData) -> None:
             f"ws://{config.ws_client_config.host}:{config.ws_client_config.port}"
         )
         self._client: Client = Client(
-            self._ws_server_url, config.ws_client_config.aiohttp_session
+            self._ws_server_url, aiohttp_session=config.ws_client_config.aiohttp_session
         )
         self._devices: dict[EUI64, WebSocketClientDevice] = {}
         self._groups: dict[int, WebSocketClientGroup] = {}

zha/websocket/__init__.py

@@ -1 +1,9 @@
 """Websocket module for Zigbee Home Automation."""
+
+from __future__ import annotations
+
+from zha.exceptions import ZHAException
+
+
+class ZHAWebSocketException(ZHAException):
+    """Exception raised by websocket errors."""

zha/websocket/client/client.py

@@ -14,7 +14,7 @@
 from async_timeout import timeout
 
 from zha.event import EventBase
-from zha.exceptions import ZHAException
+from zha.websocket import ZHAWebSocketException
 from zha.websocket.client.model.messages import Message
 from zha.websocket.server.api.model import WebSocketCommand, WebSocketCommandResponse
 
@@ -28,8 +28,8 @@ class Client(EventBase):
     def __init__(
         self,
         ws_server_url: str,
-        aiohttp_session: ClientSession | None = None,
         *args: Any,
+        aiohttp_session: ClientSession | None = None,
         **kwargs: Any,
     ) -> None:
         """Initialize the Client class."""
@@ -120,7 +120,7 @@ async def connect(self) -> None:
             )
         except client_exceptions.ClientError as err:
             _LOGGER.exception("Error connecting to server", exc_info=err)
-            raise err
+            raise ZHAWebSocketException from err
 
     async def listen_loop(self) -> None:
         """Listen to the websocket."""
@@ -132,7 +132,7 @@ async def listen_loop(self) -> None:
     async def listen(self) -> None:
         """Start listening to the websocket."""
         if not self.connected:
-            raise Exception("Not connected when start listening")  # noqa: TRY002
+            raise ZHAWebSocketException("Not connected when start listening")
 
         assert self._client
 
@@ -170,21 +170,21 @@ async def _receive_json_or_raise(self) -> dict:
         msg = await self._client.receive()
 
         if msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSED, WSMsgType.CLOSING):
-            raise Exception("Connection was closed.")  # noqa: TRY002
+            raise ZHAWebSocketException(f"Connection was closed: {msg}")
 
         if msg.type == WSMsgType.ERROR:
-            raise Exception()  # noqa: TRY002
+            raise ZHAWebSocketException(f"WS message type was ERROR: {msg}")
 
         if msg.type != WSMsgType.TEXT:
-            raise Exception(f"Received non-Text message: {msg.type}")  # noqa: TRY002
+            raise ZHAWebSocketException(f"Received non-Text message: {msg}")
 
         try:
             if len(msg.data) > SIZE_PARSE_JSON_EXECUTOR:
                 data: dict = await self._loop.run_in_executor(None, msg.json)
             else:
                 data = msg.json()
         except ValueError as err:
-            raise Exception("Received invalid JSON.") from err  # noqa: TRY002
+            raise ZHAWebSocketException(f"Received invalid JSON: {msg}") from err
 
         if _LOGGER.isEnabledFor(logging.DEBUG):
             _LOGGER.debug("Received message:\n%s\n", pprint.pformat(msg))
@@ -199,12 +199,12 @@ def _handle_incoming_message(self, msg: dict) -> None:
 
         try:
             message = Message.model_validate(msg).root
-        except Exception as err:
+        except Exception as err:  # pylint: disable=broad-except
             _LOGGER.exception("Error parsing message: %s", msg, exc_info=err)
             if msg["message_type"] == "result":
                 future = self._result_futures.get(msg["message_id"])
                 if future is not None:
-                    future.set_exception(err)
+                    future.set_exception(ZHAWebSocketException(err))
                     return
             return
 
@@ -220,9 +220,9 @@ def _handle_incoming_message(self, msg: dict) -> None:
                 return
 
             if msg["error_code"] != "zigbee_error":
-                error = ZHAException(msg["message_id"], msg["error_code"])
+                error = ZHAWebSocketException(msg["message_id"], msg["error_code"])
             else:
-                error = ZHAException(
+                error = ZHAWebSocketException(
                     msg["message_id"],
                     msg["zigbee_error_code"],
                     msg["zigbee_error_message"],
@@ -242,16 +242,17 @@ def _handle_incoming_message(self, msg: dict) -> None:
 
         try:
             self.emit(message.event_type, message)
-        except Exception as err:
+        except Exception as err:  # pylint: disable=broad-except
             _LOGGER.exception("Error handling event", exc_info=err)
+            raise ZHAWebSocketException from err
 
     async def _send_json_message(self, message: str) -> None:
         """Send a message.
 
         Raises NotConnected if client not connected.
         """
         if not self.connected:
-            raise Exception()  # noqa: TRY002
+            raise ZHAWebSocketException("Sending message failed: no active connection.")
 
         _LOGGER.debug("Publishing message:\n%s\n", pprint.pformat(message))