Add base Websocket class

Author: mj23000

homeassistant_api/rawbasewebsocket.py

@@ -0,0 +1,82 @@
+import logging
+import time
+from typing import Optional, cast
+
+from pydantic import ValidationError
+
+from homeassistant_api.errors import (
+    ReceivingError,
+    RequestError,
+)
+from homeassistant_api.models.websocket import (
+    ErrorResponse,
+    EventResponse,
+    PingResponse,
+    ResultResponse,
+)
+from homeassistant_api.utils import JSONType
+
+logger = logging.getLogger(__name__)
+
+
+class RawBaseWebsocketClient:
+    """Shared methods for Websocket clients."""
+
+    api_url: str
+    token: str
+    _id_counter: int
+    _result_responses: dict[int, Optional[ResultResponse]]
+    _event_responses: dict[int, list[EventResponse]]
+    _ping_responses: dict[int, PingResponse]
+
+    def __init__(self, api_url: str, token: str) -> None:
+        self.api_url = api_url
+        self.token = token.strip()
+
+        self._id_counter = 0
+        self._result_responses = {}  # id -> response
+        self._event_responses = {}  # id -> [response, ...]
+        self._ping_responses = {}  # id -> (sent, received)
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.api_url!r})"
+
+    def _request_id(self) -> int:
+        """Get a unique id for a message."""
+        self._id_counter += 1
+        return self._id_counter
+
+    def check_success(self, data: dict[str, JSONType]) -> None:
+        """Check if a command message was successful."""
+        try:
+            error_resp = ErrorResponse.model_validate(data)
+            raise RequestError(error_resp.error.code, error_resp.error.message)
+        except ValidationError:
+            pass
+
+    def handle_recv(self, data: dict[str, JSONType]) -> None:
+        """Handle a received message."""
+        if "id" not in data:
+            raise ReceivingError(
+                "Received a message without an id outside the auth phase."
+            )
+        self.check_success(data)
+        self.parse_response(data)
+
+    def parse_response(self, data: dict[str, JSONType]) -> None:
+        data_id = cast(int, data["id"])
+        if data.get("type") == "pong":
+            logger.info("Received pong message")
+            self._ping_responses[data_id].end = time.perf_counter_ns()
+        elif data.get("type") == "result":
+            logger.info("Received result message")
+            if data.get("success"):
+                self._result_responses[data_id] = ResultResponse.model_validate(data)
+            else:
+                error_resp = ErrorResponse.model_validate(data)
+                raise RequestError(error_resp.error.code, error_resp.error.message)
+        elif data.get("type") == "event":
+            logger.info("Received event message %s", data["event"])
+            self._event_responses[data_id].append(EventResponse.model_validate(data))
+        else:
+            raise ReceivingError(f"Received unexpected message type: {data}")

homeassistant_api/rawwebsocket.py

@@ -8,25 +8,24 @@
 
 from homeassistant_api.errors import (
     ReceivingError,
-    RequestError,
     ResponseError,
     UnauthorizedError,
 )
 from homeassistant_api.models.websocket import (
     AuthInvalid,
     AuthOk,
     AuthRequired,
-    ErrorResponse,
     EventResponse,
     PingResponse,
     ResultResponse,
 )
+from homeassistant_api.rawbasewebsocket import RawBaseWebsocketClient
 from homeassistant_api.utils import JSONType
 
 logger = logging.getLogger(__name__)
 
 
-class RawWebsocketClient:
+class RawWebsocketClient(RawBaseWebsocketClient):
     api_url: str
     token: str
     _conn: Optional[ws.ClientConnection]
@@ -36,22 +35,9 @@ def __init__(
         api_url: str,
         token: str,
     ) -> None:
-        self.api_url = api_url
-        self.token = token.strip()
+        super().__init__(api_url, token)
         self._conn = None
 
-        self._id_counter = 0
-        self._result_responses: dict[int, Optional[ResultResponse]] = (
-            {}
-        )  # id -> response
-        self._event_responses: dict[int, list[EventResponse]] = (
-            {}
-        )  # id -> [response, ...]
-        self._ping_responses: dict[int, PingResponse] = {}  # id -> (sent, received)
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}({self.api_url!r})"
-
     def __enter__(self):
         self._conn = ws.connect(self.api_url)
         self._conn.__enter__()
@@ -66,11 +52,6 @@ def __exit__(self, exc_type, exc_value, traceback):
         self._conn.__exit__(exc_type, exc_value, traceback)
         self._conn = None
 
-    def _request_id(self) -> int:
-        """Get a unique id for a message."""
-        self._id_counter += 1
-        return self._id_counter
-
     def _send(self, data: dict[str, JSONType]) -> None:
         """Send a message to the websocket server."""
         logger.debug(f"Sending message: {data}")
@@ -112,41 +93,6 @@ def send(self, type: str, include_id: bool = True, **data: Any) -> int:
             return data["id"]
         return -1  # non-command messages don't have an id
 
-    def check_success(self, data: dict[str, JSONType]) -> None:
-        """Check if a command message was successful."""
-        try:
-            error_resp = ErrorResponse.model_validate(data)
-            raise RequestError(error_resp.error.code, error_resp.error.message)
-        except ValidationError:
-            pass
-
-    def handle_recv(self, data: dict[str, JSONType]) -> None:
-        """Handle a received message."""
-        if "id" not in data:
-            raise ReceivingError(
-                "Received a message without an id outside the auth phase."
-            )
-        self.check_success(data)
-        self.parse_response(data)
-
-    def parse_response(self, data: dict[str, JSONType]) -> None:
-        data_id = cast(int, data["id"])
-        if data.get("type") == "pong":
-            logger.info("Received pong message")
-            self._ping_responses[data_id].end = time.perf_counter_ns()
-        elif data.get("type") == "result":
-            logger.info("Received result message")
-            if data.get("success"):
-                self._result_responses[data_id] = ResultResponse.model_validate(data)
-            else:
-                error_resp = ErrorResponse.model_validate(data)
-                raise RequestError(error_resp.error.code, error_resp.error.message)
-        elif data.get("type") == "event":
-            logger.info("Received event message %s", data["event"])
-            self._event_responses[data_id].append(EventResponse.model_validate(data))
-        else:
-            raise ReceivingError(f"Received unexpected message type: {data}")
-
     def recv(self, id: int) -> Union[EventResponse, ResultResponse, PingResponse]:
         """Receive a response to a message from the websocket server."""
         while True:

homeassistant_api/websocket.py

@@ -12,10 +12,9 @@
     ResultResponse,
     TemplateEvent,
 )
+from homeassistant_api.rawwebsocket import RawWebsocketClient
 from homeassistant_api.utils import JSONType, prepare_entity_id
 
-from .rawwebsocket import RawWebsocketClient
-
 logger = logging.getLogger(__name__)