python-roborock 2.9.0__py3-none-any.whl → 2.9.1__py3-none-any.whl

{python_roborock-2.9.0.dist-info → python_roborock-2.9.1.dist-info}/METADATA

@@ -1,8 +1,7 @@
 Metadata-Version: 2.3
 Name: python-roborock
-Version: 2.9.0
+Version: 2.9.1
 Summary: A package to control Roborock vacuums.
-Home-page: https://github.com/humbertogontijo/python-roborock
 License: GPL-3.0-only
 Keywords: roborock,vacuum,homeassistant
 Author: humbertogontijo

{python_roborock-2.9.0.dist-info → python_roborock-2.9.1.dist-info}/RECORD

@@ -1,29 +1,29 @@
 roborock/__init__.py,sha256=AqfxNGj4wILqL7dooJyQQd6w7o2EW02fMAF51wT9VK8,165
-roborock/api.py,sha256=zNQNf_Zmx-rW-jIbqhZm4aYnASKFFztMCtAJBUOOVzc,4270
+roborock/api.py,sha256=ahk1kUTnSefR7lPaHZJZHF-XfCFyjj6fRBVnQgruts4,4383
 roborock/cli.py,sha256=l4dOt9HZ4TZOZNp_btsNCZEDUihM94Dz173KPtUYdm4,6826
-roborock/cloud_api.py,sha256=8p_iRfXRstBW5KfdxFtya41VJjeRXbdwfpASXIqASqM,7144
+roborock/cloud_api.py,sha256=rExyj_ZLatzHFmmKlwywmUX0CQxz-a_cSjSrTaQO-cY,7109
 roborock/code_mappings.py,sha256=PTNPwWsuVb0nOkg6GpDklGMYlIB6m9uJXcMdPKVqLDY,18432
 roborock/command_cache.py,sha256=E8e7EYzYFZEmw_cIfhPYZXNpBuSJbIAT9KXnE-e5NPg,7912
 roborock/const.py,sha256=7DtkKJJ2KscgV-BKIGHsMNV_fqfVupt02LQ66r7yOyQ,2632
 roborock/containers.py,sha256=hApLhRbZlp2lMCwlecMXC-vfLMHFq7l9kOVW6rVMC3U,32050
 roborock/exceptions.py,sha256=HAe8bI78n-8-aMD_-9ahTm43xh4lc8KX9qeG5GdqtQc,1962
-roborock/local_api.py,sha256=sbgoo1ucuAu2HgdtS1IVeXpsZrJ9pBS41Qk9tB4YEtk,4375
+roborock/local_api.py,sha256=RcBp1I7sJdUfDddrd4XeJMB1OVX7L77GDeD-bPaSbIs,4652
 roborock/protocol.py,sha256=VaHCceR6zbe0_kQCnGqfkbWpqaKCq96IqAh5hgSI_Kk,12315
 roborock/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-roborock/roborock_future.py,sha256=j72NFkP1s5TXOzWkkJNr5FAec_44lRMyeSoCe37BNeM,834
+roborock/roborock_future.py,sha256=W8vm5IA1pzaKLsV2fIR9HyYFiANDLwrs8yE-OqR8ITE,1044
 roborock/roborock_message.py,sha256=5wFwaIu5fGZMG_g6PLEEjUu3Avc3Y_dEevuTeQEP9kE,5486
 roborock/roborock_typing.py,sha256=TTWU0wP-4OjdO2Qfp-uA3CySE8eupLbLJmlxTaSaHWg,23933
 roborock/util.py,sha256=CndrrzhWwMOCP1HRGsOMWtzt7EOkGS5Hr5llcAKMEUA,3946
 roborock/version_1_apis/__init__.py,sha256=BDAkBZPOtpyYUPUa6VaZd7Aaza1hbvM2fIzsjIGD9Vc,183
-roborock/version_1_apis/roborock_client_v1.py,sha256=aR9vDGRvrwwI9bw-Cr0l8AIoVypgWYN9SvNrZql1jjs,22345
-roborock/version_1_apis/roborock_local_client_v1.py,sha256=L1P8f1ldlHZtcvq6NWTcGf8fZt_DkDI4CYAY8pq8180,3564
-roborock/version_1_apis/roborock_mqtt_client_v1.py,sha256=GkQOcltXvGN-Wapsr1ZdZO3aQJHeAJ0gqDyabuSM5_Y,4270
+roborock/version_1_apis/roborock_client_v1.py,sha256=ATKmedhvwj1fqqCDKfyEb_zSHb5Dr1tfhncWpFI1AYU,22305
+roborock/version_1_apis/roborock_local_client_v1.py,sha256=SuH0yBCR4USeDg2UJdPnzqrQq8YmH39wc5ZL662IbJY,4022
+roborock/version_1_apis/roborock_mqtt_client_v1.py,sha256=qrDNry25AXAk0ZfCbQJdF2UIa5mKy19dBrNfcbBa460,4373
 roborock/version_a01_apis/__init__.py,sha256=i_WTIX0LUCG93YX3RraNaQYNYBMY-xJI7ncWMdUFQVs,111
-roborock/version_a01_apis/roborock_client_a01.py,sha256=ISXMQ7pFoohoHG2mmukdZx6J7S2sgYvvInpJGqtKGL8,7424
-roborock/version_a01_apis/roborock_mqtt_client_a01.py,sha256=fiy2LdCVYwAXo70QN5vwOM8abEHTGACmaLzaEiB8rmk,3043
+roborock/version_a01_apis/roborock_client_a01.py,sha256=84Z1uR5ZzfGOm-g4Ln4rquiL5SzNOLx0OVHUUzvjLIw,7576
+roborock/version_a01_apis/roborock_mqtt_client_a01.py,sha256=5fpUWnnPTj4yvxUwkh7ANHEJHH6RTgN3jEbVXL06blk,3180
 roborock/web_api.py,sha256=Evpql45Kh40vncbW3khK_CrPwvY1wnL6GGeBAEmG4jM,14860
-python_roborock-2.9.0.dist-info/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-python_roborock-2.9.0.dist-info/METADATA,sha256=xnk2Jk_eFimdQ9b7YuvGH4ts4olvfns_rDIAD1Dyk3A,4338
-python_roborock-2.9.0.dist-info/WHEEL,sha256=RaoafKOydTQ7I_I3JTrPCg6kUmTgtm4BornzOqyEfJ8,88
-python_roborock-2.9.0.dist-info/entry_points.txt,sha256=3V_nm8VDR5Wg8m9QlRy9R4olW2lC--6JrDjZJpQpDNI,46
-python_roborock-2.9.0.dist-info/RECORD,,
+python_roborock-2.9.1.dist-info/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+python_roborock-2.9.1.dist-info/METADATA,sha256=tuHdDtQgY5MzPxXmbAqR5zrz4Yep19wwuWVIhZ8vxiU,4276
+python_roborock-2.9.1.dist-info/WHEEL,sha256=IYZQI976HJqqOpQU6PHkJ8fb3tMNBFjg-Cn-pwAbaFM,88
+python_roborock-2.9.1.dist-info/entry_points.txt,sha256=3V_nm8VDR5Wg8m9QlRy9R4olW2lC--6JrDjZJpQpDNI,46
+python_roborock-2.9.1.dist-info/RECORD,,

{python_roborock-2.9.0.dist-info → python_roborock-2.9.1.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: poetry-core 2.0.0
+Generator: poetry-core 2.0.1
 Root-Is-Purelib: true
 Tag: py3-none-any

roborock/api.py

@@ -7,7 +7,7 @@ import base64
 import logging
 import secrets
 import time
-from collections.abc import Coroutine
+from abc import ABC, abstractmethod
 from typing import Any
 
 from .containers import (
@@ -16,24 +16,27 @@ from .containers import (
 from .exceptions import (
     RoborockTimeout,
     UnknownMethodError,
-    VacuumError,
 )
 from .roborock_future import RoborockFuture
 from .roborock_message import (
     RoborockMessage,
 )
 from .roborock_typing import RoborockCommand
-from .util import RoborockLoggerAdapter, get_next_int, get_running_loop_or_create_one
+from .util import get_next_int, get_running_loop_or_create_one
 
 _LOGGER = logging.getLogger(__name__)
 KEEPALIVE = 60
 
 
-class RoborockClient:
-    def __init__(self, endpoint: str, device_info: DeviceData, queue_timeout: int = 4) -> None:
+class RoborockClient(ABC):
+    """Roborock client base class."""
+
+    _logger: logging.LoggerAdapter
+
+    def __init__(self, device_info: DeviceData, queue_timeout: int = 4) -> None:
+        """Initialize RoborockClient."""
         self.event_loop = get_running_loop_or_create_one()
         self.device_info = device_info
-        self._endpoint = endpoint
         self._nonce = secrets.token_bytes(16)
         self._waiting_queue: dict[int, RoborockFuture] = {}
         self._last_device_msg_in = time.monotonic()
@@ -42,7 +45,6 @@ class RoborockClient:
         self._diagnostic_data: dict[str, dict[str, Any]] = {
             "misc_info": {"Nonce": base64.b64encode(self._nonce).decode("utf-8")}
         }
-        self._logger = RoborockLoggerAdapter(device_info.device.name, _LOGGER)
         self.is_available: bool = True
         self.queue_timeout = queue_timeout
 
@@ -59,17 +61,21 @@ class RoborockClient:
     def diagnostic_data(self) -> dict:
         return self._diagnostic_data
 
+    @abstractmethod
     async def async_connect(self):
-        raise NotImplementedError
+        """Connect to the Roborock device."""
 
+    @abstractmethod
     def sync_disconnect(self) -> Any:
-        raise NotImplementedError
+        """Disconnect from the Roborock device."""
 
+    @abstractmethod
     async def async_disconnect(self) -> Any:
-        raise NotImplementedError
+        """Disconnect from the Roborock device."""
 
+    @abstractmethod
     def on_message_received(self, messages: list[RoborockMessage]) -> None:
-        raise NotImplementedError
+        """Handle received incoming messages from the device."""
 
     def on_connection_lost(self, exc: Exception | None) -> None:
         self._last_disconnection = time.monotonic()
@@ -89,24 +95,22 @@ class RoborockClient:
             await self.async_disconnect()
         await self.async_connect()
 
-    async def _wait_response(self, request_id: int, queue: RoborockFuture) -> tuple[Any, VacuumError | None]:
+    async def _wait_response(self, request_id: int, queue: RoborockFuture) -> Any:
         try:
-            (response, err) = await queue.async_get(self.queue_timeout)
+            response = await queue.async_get(self.queue_timeout)
             if response == "unknown_method":
                 raise UnknownMethodError("Unknown method")
-            return response, err
+            return response
         except (asyncio.TimeoutError, asyncio.CancelledError):
             raise RoborockTimeout(f"id={request_id} Timeout after {self.queue_timeout} seconds") from None
         finally:
             self._waiting_queue.pop(request_id, None)
 
-    def _async_response(
-        self, request_id: int, protocol_id: int = 0
-    ) -> Coroutine[Any, Any, tuple[Any, VacuumError | None]]:
+    def _async_response(self, request_id: int, protocol_id: int = 0) -> Any:
         queue = RoborockFuture(protocol_id)
         if request_id in self._waiting_queue:
             new_id = get_next_int(10000, 32767)
-            _LOGGER.warning(
+            self._logger.warning(
                 "Attempting to create a future with an existing id %s (%s)... New id is %s. "
                 "Code may not function properly.",
                 request_id,
@@ -115,14 +119,16 @@ class RoborockClient:
             )
             request_id = new_id
         self._waiting_queue[request_id] = queue
-        return self._wait_response(request_id, queue)
+        return asyncio.ensure_future(self._wait_response(request_id, queue))
 
+    @abstractmethod
     async def send_message(self, roborock_message: RoborockMessage):
-        raise NotImplementedError
+        """Send a message to the Roborock device."""
 
+    @abstractmethod
     async def _send_command(
         self,
         method: RoborockCommand | str,
         params: list | dict | int | None = None,
     ):
-        raise NotImplementedError
+        """Send a command to the Roborock device."""

roborock/cloud_api.py

@@ -1,12 +1,10 @@
 from __future__ import annotations
 
-import asyncio
-import base64
 import logging
 import threading
-import typing
 import uuid
-from asyncio import Lock, Task
+from abc import ABC
+from asyncio import Lock
 from typing import Any
 from urllib.parse import urlparse
 
@@ -15,31 +13,52 @@ import paho.mqtt.client as mqtt
 from .api import KEEPALIVE, RoborockClient
 from .containers import DeviceData, UserData
 from .exceptions import RoborockException, VacuumError
-from .protocol import MessageParser, Utils, md5hex
+from .protocol import MessageParser, md5hex
 from .roborock_future import RoborockFuture
-from .roborock_message import RoborockMessage
-from .roborock_typing import RoborockCommand
-from .util import RoborockLoggerAdapter
 
-if typing.TYPE_CHECKING:
-    pass
 _LOGGER = logging.getLogger(__name__)
 CONNECT_REQUEST_ID = 0
 DISCONNECT_REQUEST_ID = 1
 
 
-class RoborockMqttClient(RoborockClient, mqtt.Client):
+class _Mqtt(mqtt.Client):
+    """Internal MQTT client.
+
+    This is a subclass of the Paho MQTT client that adds some additional functionality
+    for error cases where things get stuck.
+    """
+
     _thread: threading.Thread
     _client_id: str
 
+    def __init__(self) -> None:
+        """Initialize the MQTT client."""
+        super().__init__(protocol=mqtt.MQTTv5)
+        self.reset_client_id()
+
+    def reset_client_id(self):
+        """Generate a new client id to make a new session when reconnecting."""
+        self._client_id = mqtt.base62(uuid.uuid4().int, padding=22)
+
+    def maybe_restart_loop(self) -> None:
+        """Ensure that the MQTT loop is running in case it previously exited."""
+        if not self._thread or not self._thread.is_alive():
+            if self._thread:
+                _LOGGER.info("Stopping mqtt loop")
+                super().loop_stop()
+            _LOGGER.info("Starting mqtt loop")
+            super().loop_start()
+
+
+class RoborockMqttClient(RoborockClient, ABC):
+    """Roborock MQTT client base class."""
+
     def __init__(self, user_data: UserData, device_info: DeviceData, queue_timeout: int = 10) -> None:
+        """Initialize the Roborock MQTT client."""
         rriot = user_data.rriot
         if rriot is None:
             raise RoborockException("Got no rriot data from user_data")
-        endpoint = base64.b64encode(Utils.md5(rriot.k.encode())[8:14]).decode()
-        RoborockClient.__init__(self, endpoint, device_info, queue_timeout)
-        mqtt.Client.__init__(self, protocol=mqtt.MQTTv5)
-        self._logger = RoborockLoggerAdapter(device_info.device.name, _LOGGER)
+        RoborockClient.__init__(self, device_info, queue_timeout)
         self._mqtt_user = rriot.u
         self._hashed_user = md5hex(self._mqtt_user + ":" + rriot.k)[2:10]
         url = urlparse(rriot.r.m)
@@ -48,39 +67,43 @@ class RoborockMqttClient(RoborockClient, mqtt.Client):
         self._mqtt_host = str(url.hostname)
         self._mqtt_port = url.port
         self._mqtt_ssl = url.scheme == "ssl"
+
+        self._mqtt_client = _Mqtt()
+        self._mqtt_client.on_connect = self._mqtt_on_connect
+        self._mqtt_client.on_message = self._mqtt_on_message
+        self._mqtt_client.on_disconnect = self._mqtt_on_disconnect
         if self._mqtt_ssl:
-            super().tls_set()
+            self._mqtt_client.tls_set()
+
         self._mqtt_password = rriot.s
         self._hashed_password = md5hex(self._mqtt_password + ":" + rriot.k)[16:]
-        super().username_pw_set(self._hashed_user, self._hashed_password)
-        self._endpoint = base64.b64encode(Utils.md5(rriot.k.encode())[8:14]).decode()
+        self._mqtt_client.username_pw_set(self._hashed_user, self._hashed_password)
         self._waiting_queue: dict[int, RoborockFuture] = {}
         self._mutex = Lock()
-        self.update_client_id()
 
-    def on_connect(self, *args, **kwargs):
+    def _mqtt_on_connect(self, *args, **kwargs):
         _, __, ___, rc, ____ = args
         connection_queue = self._waiting_queue.get(CONNECT_REQUEST_ID)
         if rc != mqtt.MQTT_ERR_SUCCESS:
             message = f"Failed to connect ({mqtt.error_string(rc)})"
             self._logger.error(message)
             if connection_queue:
-                connection_queue.resolve((None, VacuumError(message)))
+                connection_queue.set_exception(VacuumError(message))
             return
         self._logger.info(f"Connected to mqtt {self._mqtt_host}:{self._mqtt_port}")
         topic = f"rr/m/o/{self._mqtt_user}/{self._hashed_user}/{self.device_info.device.duid}"
-        (result, mid) = self.subscribe(topic)
+        (result, mid) = self._mqtt_client.subscribe(topic)
         if result != 0:
             message = f"Failed to subscribe ({mqtt.error_string(rc)})"
             self._logger.error(message)
             if connection_queue:
-                connection_queue.resolve((None, VacuumError(message)))
+                connection_queue.set_exception(VacuumError(message))
             return
         self._logger.info(f"Subscribed to topic {topic}")
         if connection_queue:
-            connection_queue.resolve((True, None))
+            connection_queue.set_result(True)
 
-    def on_message(self, *args, **kwargs):
+    def _mqtt_on_message(self, *args, **kwargs):
         client, __, msg = args
         try:
             messages, _ = MessageParser.parse(msg.payload, local_key=self.device_info.device.local_key)
@@ -88,93 +111,74 @@ class RoborockMqttClient(RoborockClient, mqtt.Client):
         except Exception as ex:
             self._logger.exception(ex)
 
-    def on_disconnect(self, *args, **kwargs):
+    def _mqtt_on_disconnect(self, *args, **kwargs):
         _, __, rc, ___ = args
         try:
             exc = RoborockException(mqtt.error_string(rc)) if rc != mqtt.MQTT_ERR_SUCCESS else None
             super().on_connection_lost(exc)
             if rc == mqtt.MQTT_ERR_PROTOCOL:
-                self.update_client_id()
+                self._mqtt_client.reset_client_id()
             connection_queue = self._waiting_queue.get(DISCONNECT_REQUEST_ID)
             if connection_queue:
-                connection_queue.resolve((True, None))
+                connection_queue.set_result(True)
         except Exception as ex:
             self._logger.exception(ex)
 
-    def update_client_id(self):
-        self._client_id = mqtt.base62(uuid.uuid4().int, padding=22)
-
-    def sync_stop_loop(self) -> None:
-        if self._thread:
-            self._logger.info("Stopping mqtt loop")
-            super().loop_stop()
-
-    def sync_start_loop(self) -> None:
-        if not self._thread or not self._thread.is_alive():
-            self.sync_stop_loop()
-            self._logger.info("Starting mqtt loop")
-            super().loop_start()
+    def is_connected(self) -> bool:
+        """Check if the mqtt client is connected."""
+        return self._mqtt_client.is_connected()
 
-    def sync_disconnect(self) -> tuple[bool, Task[tuple[Any, VacuumError | None]] | None]:
+    def sync_disconnect(self) -> Any:
         if not self.is_connected():
-            return False, None
+            return None
 
         self._logger.info("Disconnecting from mqtt")
-        disconnected_future = asyncio.ensure_future(self._async_response(DISCONNECT_REQUEST_ID))
-        rc = super().disconnect()
+        disconnected_future = self._async_response(DISCONNECT_REQUEST_ID)
+        rc = self._mqtt_client.disconnect()
 
         if rc == mqtt.MQTT_ERR_NO_CONN:
             disconnected_future.cancel()
-            return False, None
+            return None
 
         if rc != mqtt.MQTT_ERR_SUCCESS:
             disconnected_future.cancel()
             raise RoborockException(f"Failed to disconnect ({mqtt.error_string(rc)})")
 
-        return True, disconnected_future
+        return disconnected_future
 
-    def sync_connect(self) -> tuple[bool, Task[tuple[Any, VacuumError | None]] | None]:
+    def sync_connect(self) -> Any:
         if self.is_connected():
-            self.sync_start_loop()
-            return False, None
+            self._mqtt_client.maybe_restart_loop()
+            return None
 
         if self._mqtt_port is None or self._mqtt_host is None:
             raise RoborockException("Mqtt information was not entered. Cannot connect.")
 
         self._logger.debug("Connecting to mqtt")
-        connected_future = asyncio.ensure_future(self._async_response(CONNECT_REQUEST_ID))
-        super().connect(host=self._mqtt_host, port=self._mqtt_port, keepalive=KEEPALIVE)
-
-        self.sync_start_loop()
-        return True, connected_future
+        connected_future = self._async_response(CONNECT_REQUEST_ID)
+        self._mqtt_client.connect(host=self._mqtt_host, port=self._mqtt_port, keepalive=KEEPALIVE)
+        self._mqtt_client.maybe_restart_loop()
+        return connected_future
 
     async def async_disconnect(self) -> None:
         async with self._mutex:
-            (disconnecting, disconnected_future) = self.sync_disconnect()
-            if disconnecting and disconnected_future:
-                (_, err) = await disconnected_future
-                if err:
+            if disconnected_future := self.sync_disconnect():
+                try:
+                    await disconnected_future
+                except VacuumError as err:
                     raise RoborockException(err) from err
 
     async def async_connect(self) -> None:
         async with self._mutex:
-            (connecting, connected_future) = self.sync_connect()
-            if connecting and connected_future:
-                (_, err) = await connected_future
-                if err:
+            if connected_future := self.sync_connect():
+                try:
+                    await connected_future
+                except VacuumError as err:
                     raise RoborockException(err) from err
 
     def _send_msg_raw(self, msg: bytes) -> None:
-        info = self.publish(f"rr/m/i/{self._mqtt_user}/{self._hashed_user}/{self.device_info.device.duid}", msg)
+        info = self._mqtt_client.publish(
+            f"rr/m/i/{self._mqtt_user}/{self._hashed_user}/{self.device_info.device.duid}", msg
+        )
         if info.rc != mqtt.MQTT_ERR_SUCCESS:
             raise RoborockException(f"Failed to publish ({mqtt.error_string(info.rc)})")
-
-    async def send_message(self, roborock_message: RoborockMessage):
-        raise NotImplementedError
-
-    async def _send_command(
-        self,
-        method: RoborockCommand | str,
-        params: list | dict | int | None = None,
-    ):
-        raise NotImplementedError

roborock/local_api.py

@@ -2,7 +2,10 @@ from __future__ import annotations
 
 import asyncio
 import logging
+from abc import ABC
 from asyncio import Lock, TimerHandle, Transport
+from collections.abc import Callable
+from dataclasses import dataclass
 
 import async_timeout
 
@@ -11,14 +14,23 @@ from .api import RoborockClient
 from .exceptions import RoborockConnectionException, RoborockException
 from .protocol import MessageParser
 from .roborock_message import RoborockMessage, RoborockMessageProtocol
-from .roborock_typing import RoborockCommand
-from .util import RoborockLoggerAdapter
 
 _LOGGER = logging.getLogger(__name__)
 
 
-class RoborockLocalClient(RoborockClient, asyncio.Protocol):
+@dataclass
+class _LocalProtocol(asyncio.Protocol):
+    """Callbacks for the Roborock local client transport."""
+
+    messages_cb: Callable[[bytes], None]
+    connection_lost_cb: Callable[[Exception | None], None]
+
+
+class RoborockLocalClient(RoborockClient, ABC):
+    """Roborock local client base class."""
+
     def __init__(self, device_data: DeviceData, queue_timeout: int = 4):
+        """Initialize the Roborock local client."""
         if device_data.host is None:
             raise RoborockException("Host is required")
         self.host = device_data.host
@@ -28,17 +40,19 @@ class RoborockLocalClient(RoborockClient, asyncio.Protocol):
         self.transport: Transport | None = None
         self._mutex = Lock()
         self.keep_alive_task: TimerHandle | None = None
-        self._logger = RoborockLoggerAdapter(device_data.device.name, _LOGGER)
-        RoborockClient.__init__(self, "abc", device_data, queue_timeout)
+        RoborockClient.__init__(self, device_data, queue_timeout)
+        self._local_protocol = _LocalProtocol(self._data_received, self._connection_lost)
 
-    def data_received(self, message):
+    def _data_received(self, message):
+        """Called when data is received from the transport."""
         if self.remaining:
             message = self.remaining + message
             self.remaining = b""
         parser_msg, self.remaining = MessageParser.parse(message, local_key=self.device_info.device.local_key)
         self.on_message_received(parser_msg)
 
-    def connection_lost(self, exc: Exception | None):
+    def _connection_lost(self, exc: Exception | None):
+        """Called when the transport connection is lost."""
         self.sync_disconnect()
         self.on_connection_lost(exc)
 
@@ -61,7 +75,7 @@ class RoborockLocalClient(RoborockClient, asyncio.Protocol):
                     async with async_timeout.timeout(self.queue_timeout):
                         self._logger.debug(f"Connecting to {self.host}")
                         self.transport, _ = await self.event_loop.create_connection(  # type: ignore
-                            lambda: self, self.host, 58867
+                            lambda: self._local_protocol, self.host, 58867
                         )
                         self._logger.info(f"Connected to {self.host}")
                         should_ping = True
@@ -107,13 +121,6 @@ class RoborockLocalClient(RoborockClient, asyncio.Protocol):
             )
         )
 
-    async def _send_command(
-        self,
-        method: RoborockCommand | str,
-        params: list | dict | int | None = None,
-    ):
-        raise NotImplementedError
-
     def _send_msg_raw(self, data: bytes):
         try:
             if not self.transport:

roborock/roborock_future.py

@@ -14,12 +14,19 @@ class RoborockFuture:
         self.fut: Future = Future()
         self.loop = self.fut.get_loop()
 
-    def _resolve(self, item: tuple[Any, VacuumError | None]) -> None:
+    def _set_result(self, item: Any) -> None:
         if not self.fut.cancelled():
             self.fut.set_result(item)
 
-    def resolve(self, item: tuple[Any, VacuumError | None]) -> None:
-        self.loop.call_soon_threadsafe(self._resolve, item)
+    def set_result(self, item: Any) -> None:
+        self.loop.call_soon_threadsafe(self._set_result, item)
+
+    def _set_exception(self, exc: VacuumError) -> None:
+        if not self.fut.cancelled():
+            self.fut.set_exception(exc)
+
+    def set_exception(self, exc: VacuumError) -> None:
+        self.loop.call_soon_threadsafe(self._set_exception, exc)
 
     async def async_get(self, timeout: float | int) -> tuple[Any, VacuumError | None]:
         try:

roborock/version_1_apis/roborock_client_v1.py

@@ -4,6 +4,7 @@ import json
 import math
 import struct
 import time
+from abc import ABC
 from collections.abc import Callable, Coroutine
 from typing import Any, TypeVar, final
 
@@ -142,19 +143,22 @@ class ListenerModel:
     cache: dict[CacheableAttribute, AttributeCache]
 
 
-class RoborockClientV1(RoborockClient):
+class RoborockClientV1(RoborockClient, ABC):
+    """Roborock client base class for version 1 devices."""
+
     _listeners: dict[str, ListenerModel] = {}
 
-    def __init__(self, device_info: DeviceData, logger, endpoint: str):
-        super().__init__(endpoint, device_info)
+    def __init__(self, device_info: DeviceData, endpoint: str):
+        """Initializes the Roborock client."""
+        super().__init__(device_info)
         self._status_type: type[Status] = ModelStatus.get(device_info.model, S7MaxVStatus)
-        self._logger = logger
         self.cache: dict[CacheableAttribute, AttributeCache] = {
             cacheable_attribute: AttributeCache(attr, self) for cacheable_attribute, attr in get_cache_map().items()
         }
         if device_info.device.duid not in self._listeners:
             self._listeners[device_info.device.duid] = ListenerModel({}, self.cache)
         self.listener_model = self._listeners[device_info.device.duid]
+        self._endpoint = endpoint
 
     def release(self):
         super().release()
@@ -378,20 +382,17 @@ class RoborockClientV1(RoborockClient):
                             if queue and queue.protocol == protocol:
                                 error = data_point_response.get("error")
                                 if error:
-                                    queue.resolve(
-                                        (
-                                            None,
-                                            VacuumError(
-                                                error.get("code"),
-                                                error.get("message"),
-                                            ),
-                                        )
+                                    queue.set_exception(
+                                        VacuumError(
+                                            error.get("code"),
+                                            error.get("message"),
+                                        ),
                                     )
                                 else:
                                     result = data_point_response.get("result")
                                     if isinstance(result, list) and len(result) == 1:
                                         result = result[0]
-                                    queue.resolve((result, None))
+                                    queue.set_result(result)
                             else:
                                 self._logger.debug("Received response for unknown request id %s", request_id)
                         else:
@@ -451,13 +452,13 @@ class RoborockClientV1(RoborockClient):
                         if queue:
                             if isinstance(decompressed, list):
                                 decompressed = decompressed[0]
-                            queue.resolve((decompressed, None))
+                            queue.set_result(decompressed)
                         else:
                             self._logger.debug("Received response for unknown request id %s", request_id)
                 else:
                     queue = self._waiting_queue.get(data.seq)
                     if queue:
-                        queue.resolve((data.payload, None))
+                        queue.set_result(data.payload)
                     else:
                         self._logger.debug("Received response for unknown request id %s", data.seq)
         except Exception as ex:

roborock/version_1_apis/roborock_local_client_v1.py

@@ -1,17 +1,25 @@
-import asyncio
+import logging
 
 from roborock.local_api import RoborockLocalClient
 
 from .. import CommandVacuumError, DeviceData, RoborockCommand, RoborockException
+from ..exceptions import VacuumError
 from ..protocol import MessageParser
 from ..roborock_message import MessageRetry, RoborockMessage, RoborockMessageProtocol
+from ..util import RoborockLoggerAdapter
 from .roborock_client_v1 import COMMANDS_SECURED, RoborockClientV1
 
+_LOGGER = logging.getLogger(__name__)
+
 
 class RoborockLocalClientV1(RoborockLocalClient, RoborockClientV1):
+    """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, queue_timeout)
-        RoborockClientV1.__init__(self, device_data, self._logger, "abc")
+        RoborockClientV1.__init__(self, device_data, "abc")
+        self._logger = RoborockLoggerAdapter(device_data.device.name, _LOGGER)
 
     def build_roborock_message(
         self, method: RoborockCommand | str, params: list | dict | int | None = None
@@ -53,16 +61,21 @@ class RoborockLocalClientV1(RoborockLocalClient, RoborockClientV1):
         if method:
             self._logger.debug(f"id={request_id} Requesting method {method} with {params}")
         # Send the command to the Roborock device
-        async_response = asyncio.ensure_future(self._async_response(request_id, response_protocol))
+        async_response = self._async_response(request_id, response_protocol)
         self._send_msg_raw(msg)
-        (response, err) = await async_response
-        self._diagnostic_data[method if method is not None else "unknown"] = {
+        diagnostic_key = method if method is not None else "unknown"
+        try:
+            response = await async_response
+        except VacuumError as err:
+            self._diagnostic_data[diagnostic_key] = {
+                "params": roborock_message.get_params(),
+                "error": err,
+            }
+            raise CommandVacuumError(method, err) from err
+        self._diagnostic_data[diagnostic_key] = {
             "params": roborock_message.get_params(),
             "response": response,
-            "error": err,
         }
-        if err:
-            raise CommandVacuumError(method, err) from err
         if roborock_message.protocol == RoborockMessageProtocol.GENERAL_REQUEST:
             self._logger.debug(f"id={request_id} Response from method {roborock_message.get_method()}: {response}")
         if response == "retry":

roborock/version_1_apis/roborock_mqtt_client_v1.py

@@ -1,7 +1,6 @@
-import asyncio
 import base64
+import logging
 
-import paho.mqtt.client as mqtt
 from vacuum_map_parser_base.config.color import ColorsPalette
 from vacuum_map_parser_base.config.image_config import ImageConfig
 from vacuum_map_parser_base.config.size import Sizes
@@ -10,30 +9,32 @@ from vacuum_map_parser_roborock.map_data_parser import RoborockMapDataParser
 from roborock.cloud_api import RoborockMqttClient
 
 from ..containers import DeviceData, UserData
-from ..exceptions import CommandVacuumError, RoborockException
+from ..exceptions import CommandVacuumError, RoborockException, VacuumError
 from ..protocol import MessageParser, Utils
 from ..roborock_message import (
     RoborockMessage,
     RoborockMessageProtocol,
 )
 from ..roborock_typing import RoborockCommand
+from ..util import RoborockLoggerAdapter
 from .roborock_client_v1 import COMMANDS_SECURED, CUSTOM_COMMANDS, RoborockClientV1
 
+_LOGGER = logging.getLogger(__name__)
+
 
 class RoborockMqttClientV1(RoborockMqttClient, RoborockClientV1):
+    """Roborock mqtt client for v1 devices."""
+
     def __init__(self, user_data: UserData, device_info: DeviceData, queue_timeout: int = 10) -> None:
+        """Initialize the Roborock mqtt client."""
         rriot = user_data.rriot
         if rriot is None:
             raise RoborockException("Got no rriot data from user_data")
         endpoint = base64.b64encode(Utils.md5(rriot.k.encode())[8:14]).decode()
 
         RoborockMqttClient.__init__(self, user_data, device_info, queue_timeout)
-        RoborockClientV1.__init__(self, device_info, self._logger, endpoint)
-
-    def _send_msg_raw(self, msg: bytes) -> None:
-        info = self.publish(f"rr/m/i/{self._mqtt_user}/{self._hashed_user}/{self.device_info.device.duid}", msg)
-        if info.rc != mqtt.MQTT_ERR_SUCCESS:
-            raise RoborockException(f"Failed to publish ({mqtt.error_string(info.rc)})")
+        RoborockClientV1.__init__(self, device_info, endpoint)
+        self._logger = RoborockLoggerAdapter(device_info.device.name, _LOGGER)
 
     async def send_message(self, roborock_message: RoborockMessage):
         await self.validate_connection()
@@ -49,16 +50,21 @@ class RoborockMqttClientV1(RoborockMqttClient, RoborockClientV1):
         local_key = self.device_info.device.local_key
         msg = MessageParser.build(roborock_message, local_key, False)
         self._logger.debug(f"id={request_id} Requesting method {method} with {params}")
-        async_response = asyncio.ensure_future(self._async_response(request_id, response_protocol))
+        async_response = self._async_response(request_id, response_protocol)
         self._send_msg_raw(msg)
-        (response, err) = await async_response
-        self._diagnostic_data[method if method is not None else "unknown"] = {
+        diagnostic_key = method if method is not None else "unknown"
+        try:
+            response = await async_response
+        except VacuumError as err:
+            self._diagnostic_data[diagnostic_key] = {
+                "params": roborock_message.get_params(),
+                "error": err,
+            }
+            raise CommandVacuumError(method, err) from err
+        self._diagnostic_data[diagnostic_key] = {
             "params": roborock_message.get_params(),
             "response": response,
-            "error": err,
         }
-        if err:
-            raise CommandVacuumError(method, err) from err
         if response_protocol == RoborockMessageProtocol.MAP_RESPONSE:
             self._logger.debug(f"id={request_id} Response from {method}: {len(response)} bytes")
         else:

roborock/version_a01_apis/roborock_client_a01.py

@@ -1,6 +1,8 @@
 import dataclasses
 import json
+import logging
 import typing
+from abc import ABC, abstractmethod
 from collections.abc import Callable
 from datetime import time
 
@@ -38,6 +40,8 @@ from roborock.roborock_message import (
     RoborockZeoProtocol,
 )
 
+_LOGGER = logging.getLogger(__name__)
+
 
 @dataclasses.dataclass
 class A01ProtocolCacheEntry:
@@ -101,9 +105,12 @@ zeo_data_protocol_entries = {
 }
 
 
-class RoborockClientA01(RoborockClient):
-    def __init__(self, endpoint: str, device_info: DeviceData, category: RoborockCategory, queue_timeout: int = 4):
-        super().__init__(endpoint, device_info, queue_timeout)
+class RoborockClientA01(RoborockClient, ABC):
+    """Roborock client base class for A01 devices."""
+
+    def __init__(self, device_info: DeviceData, category: RoborockCategory, queue_timeout: int = 4):
+        """Initialize the Roborock client."""
+        super().__init__(device_info, queue_timeout)
         self.category = category
 
     def on_message_received(self, messages: list[RoborockMessage]) -> None:
@@ -135,10 +142,10 @@ class RoborockClientA01(RoborockClient):
                         converted_response = entries[data_point_protocol].post_process_fn(data_point)
                         queue = self._waiting_queue.get(int(data_point_number))
                         if queue and queue.protocol == protocol:
-                            queue.resolve((converted_response, None))
+                            queue.set_result(converted_response)
 
+    @abstractmethod
     async def update_values(
         self, dyad_data_protocols: list[RoborockDyadDataProtocol | RoborockZeoProtocol]
     ) -> dict[RoborockDyadDataProtocol | RoborockZeoProtocol, typing.Any]:
         """This should handle updating for each given protocol."""
-        raise NotImplementedError

roborock/version_a01_apis/roborock_mqtt_client_a01.py

@@ -1,6 +1,6 @@
 import asyncio
-import base64
 import json
+import logging
 import typing
 
 from Crypto.Cipher import AES
@@ -9,7 +9,7 @@ from Crypto.Util.Padding import pad, unpad
 from roborock.cloud_api import RoborockMqttClient
 from roborock.containers import DeviceData, RoborockCategory, UserData
 from roborock.exceptions import RoborockException
-from roborock.protocol import MessageParser, Utils
+from roborock.protocol import MessageParser
 from roborock.roborock_message import (
     RoborockDyadDataProtocol,
     RoborockMessage,
@@ -17,20 +17,26 @@ from roborock.roborock_message import (
     RoborockZeoProtocol,
 )
 
+from ..util import RoborockLoggerAdapter
 from .roborock_client_a01 import RoborockClientA01
 
+_LOGGER = logging.getLogger(__name__)
+
 
 class RoborockMqttClientA01(RoborockMqttClient, RoborockClientA01):
+    """Roborock mqtt client for A01 devices."""
+
     def __init__(
         self, user_data: UserData, device_info: DeviceData, category: RoborockCategory, queue_timeout: int = 10
     ) -> None:
+        """Initialize the Roborock mqtt client."""
         rriot = user_data.rriot
         if rriot is None:
             raise RoborockException("Got no rriot data from user_data")
-        endpoint = base64.b64encode(Utils.md5(rriot.k.encode())[8:14]).decode()
 
         RoborockMqttClient.__init__(self, user_data, device_info, queue_timeout)
-        RoborockClientA01.__init__(self, endpoint, device_info, category, queue_timeout)
+        RoborockClientA01.__init__(self, device_info, category, queue_timeout)
+        self._logger = RoborockLoggerAdapter(device_info.device.name, _LOGGER)
 
     async def send_message(self, roborock_message: RoborockMessage):
         await self.validate_connection()
@@ -43,7 +49,7 @@ class RoborockMqttClientA01(RoborockMqttClient, RoborockClientA01):
         futures = []
         if "10000" in payload["dps"]:
             for dps in json.loads(payload["dps"]["10000"]):
-                futures.append(asyncio.ensure_future(self._async_response(dps, response_protocol)))
+                futures.append(self._async_response(dps, response_protocol))
         self._send_msg_raw(m)
         responses = await asyncio.gather(*futures, return_exceptions=True)
         dps_responses: dict[int, typing.Any] = {}
@@ -54,7 +60,7 @@ class RoborockMqttClientA01(RoborockMqttClient, RoborockClientA01):
                     self._logger.warning("Timed out get req for %s after %s s", dps, self.queue_timeout)
                     dps_responses[dps] = None
                 else:
-                    dps_responses[dps] = response[0]
+                    dps_responses[dps] = response
         return dps_responses
 
     async def update_values(