fix: Merge the local api with the local v1 api (#438)

Author: allenporter

roborock/api.py

@@ -116,7 +116,3 @@ def _async_response(self, request_id: int, protocol_id: int = 0) -> Any:
             request_id = new_id
         self._waiting_queue[request_id] = queue
         return asyncio.ensure_future(self._wait_response(request_id, queue))
-
-    @abstractmethod
-    async def send_message(self, roborock_message: RoborockMessage):
-        """Send a message to the Roborock device."""

roborock/local_api.py

@@ -1,9 +1,15 @@
+import asyncio
 import logging
+from asyncio import Lock, TimerHandle, Transport, get_running_loop
+from collections.abc import Callable
+from dataclasses import dataclass
 
-from roborock.local_api import RoborockLocalClient
+import async_timeout
 
-from .. import CommandVacuumError, DeviceData, RoborockCommand, RoborockException
-from ..exceptions import VacuumError
+from .. import CommandVacuumError, DeviceData, RoborockCommand
+from ..api import RoborockClient
+from ..exceptions import RoborockConnectionException, RoborockException, VacuumError
+from ..protocol import Decoder, Encoder, create_local_decoder, create_local_encoder
 from ..protocols.v1_protocol import encode_local_payload
 from ..roborock_message import RoborockMessage, RoborockMessageProtocol
 from ..util import RoborockLoggerAdapter
@@ -12,16 +18,129 @@
 _LOGGER = logging.getLogger(__name__)
 
 
-class RoborockLocalClientV1(RoborockLocalClient, RoborockClientV1):
+@dataclass
+class _LocalProtocol(asyncio.Protocol):
+    """Callbacks for the Roborock local client transport."""
+
+    messages_cb: Callable[[bytes], None]
+    connection_lost_cb: Callable[[Exception | None], None]
+
+    def data_received(self, bytes) -> None:
+        """Called when data is received from the transport."""
+        self.messages_cb(bytes)
+
+    def connection_lost(self, exc: Exception | None) -> None:
+        """Called when the transport connection is lost."""
+        self.connection_lost_cb(exc)
+
+
+class RoborockLocalClientV1(RoborockClientV1, RoborockClient):
     """Roborock local client for v1 devices."""
 
     def __init__(self, device_data: DeviceData, queue_timeout: int = 4):
         """Initialize the Roborock local client."""
-        RoborockLocalClient.__init__(self, device_data)
+        if device_data.host is None:
+            raise RoborockException("Host is required")
+        self.host = device_data.host
+        self._batch_structs: list[RoborockMessage] = []
+        self._executing = False
+        self.transport: Transport | None = None
+        self._mutex = Lock()
+        self.keep_alive_task: TimerHandle | None = None
         RoborockClientV1.__init__(self, device_data, "abc")
+        RoborockClient.__init__(self, device_data)
+        self._local_protocol = _LocalProtocol(self._data_received, self._connection_lost)
+        self._encoder: Encoder = create_local_encoder(device_data.device.local_key)
+        self._decoder: Decoder = create_local_decoder(device_data.device.local_key)
         self.queue_timeout = queue_timeout
         self._logger = RoborockLoggerAdapter(device_data.device.name, _LOGGER)
 
+    def _data_received(self, message):
+        """Called when data is received from the transport."""
+        parsed_msg = self._decoder(message)
+        self.on_message_received(parsed_msg)
+
+    def _connection_lost(self, exc: Exception | None):
+        """Called when the transport connection is lost."""
+        self._sync_disconnect()
+        self.on_connection_lost(exc)
+
+    def is_connected(self):
+        return self.transport and self.transport.is_reading()
+
+    async def keep_alive_func(self, _=None):
+        try:
+            await self.ping()
+        except RoborockException:
+            pass
+        loop = asyncio.get_running_loop()
+        self.keep_alive_task = loop.call_later(10, lambda: asyncio.create_task(self.keep_alive_func()))
+
+    async def async_connect(self) -> None:
+        should_ping = False
+        async with self._mutex:
+            try:
+                if not self.is_connected():
+                    self._sync_disconnect()
+                    async with async_timeout.timeout(self.queue_timeout):
+                        self._logger.debug(f"Connecting to {self.host}")
+                        loop = get_running_loop()
+                        self.transport, _ = await loop.create_connection(  # type: ignore
+                            lambda: self._local_protocol, self.host, 58867
+                        )
+                        self._logger.info(f"Connected to {self.host}")
+                        should_ping = True
+            except BaseException as e:
+                raise RoborockConnectionException(f"Failed connecting to {self.host}") from e
+        if should_ping:
+            await self.hello()
+            await self.keep_alive_func()
+
+    def _sync_disconnect(self) -> None:
+        loop = asyncio.get_running_loop()
+        if self.transport and loop.is_running():
+            self._logger.debug(f"Disconnecting from {self.host}")
+            self.transport.close()
+        if self.keep_alive_task:
+            self.keep_alive_task.cancel()
+
+    async def async_disconnect(self) -> None:
+        async with self._mutex:
+            self._sync_disconnect()
+
+    async def hello(self):
+        request_id = 1
+        protocol = RoborockMessageProtocol.HELLO_REQUEST
+        try:
+            return await self._send_message(
+                RoborockMessage(
+                    protocol=protocol,
+                    seq=request_id,
+                    random=22,
+                )
+            )
+        except Exception as e:
+            self._logger.error(e)
+
+    async def ping(self) -> None:
+        request_id = 2
+        protocol = RoborockMessageProtocol.PING_REQUEST
+        return await self._send_message(
+            RoborockMessage(
+                protocol=protocol,
+                seq=request_id,
+                random=23,
+            )
+        )
+
+    def _send_msg_raw(self, data: bytes):
+        try:
+            if not self.transport:
+                raise RoborockException("Can not send message without connection")
+            self.transport.write(data)
+        except Exception as e:
+            raise RoborockException(e) from e
+
     async def _send_command(
         self,
         method: RoborockCommand | str,
@@ -32,9 +151,9 @@ async def _send_command(
 
         roborock_message = encode_local_payload(method, params)
         self._logger.debug("Building message id %s for method %s", roborock_message.get_request_id(), method)
-        return await self.send_message(roborock_message)
+        return await self._send_message(roborock_message)
 
-    async def send_message(self, roborock_message: RoborockMessage):
+    async def _send_message(self, roborock_message: RoborockMessage):
         await self.validate_connection()
         method = roborock_message.get_method()
         params = roborock_message.get_params()

roborock/version_1_apis/roborock_local_client_v1.py

@@ -39,7 +39,7 @@ def __init__(
         self.queue_timeout = queue_timeout
         self._logger = RoborockLoggerAdapter(device_info.device.name, _LOGGER)
 
-    async def send_message(self, roborock_message: RoborockMessage):
+    async def _send_message(self, roborock_message: RoborockMessage):
         await self.validate_connection()
         response_protocol = RoborockMessageProtocol.RPC_RESPONSE
 
@@ -67,11 +67,11 @@ async def update_values(
         message = encode_mqtt_payload(
             {RoborockDyadDataProtocol.ID_QUERY: str([int(protocol) for protocol in dyad_data_protocols])}
         )
-        return await self.send_message(message)
+        return await self._send_message(message)
 
     async def set_value(
         self, protocol: RoborockDyadDataProtocol | RoborockZeoProtocol, value: typing.Any
     ) -> dict[int, typing.Any]:
         """Set a value for a specific protocol on the A01 device."""
         message = encode_mqtt_payload({protocol: value})
-        return await self.send_message(message)
+        return await self._send_message(message)

roborock/version_a01_apis/roborock_mqtt_client_a01.py

@@ -307,7 +307,7 @@ def handle_write(data: bytes) -> None:
 
         return (mock_transport, "proto")
 
-    with patch("roborock.local_api.get_running_loop") as mock_loop:
+    with patch("roborock.version_1_apis.roborock_local_client_v1.get_running_loop") as mock_loop:
         mock_loop.return_value.create_connection.side_effect = create_connection
         yield