refactor: simplify future usage within the api clients

Author: allenporter

roborock/api.py

@@ -7,7 +7,6 @@
 import logging
 import secrets
 import time
-from collections.abc import Coroutine
 from typing import Any
 
 from .containers import (
@@ -16,7 +15,6 @@
 from .exceptions import (
     RoborockTimeout,
     UnknownMethodError,
-    VacuumError,
 )
 from .roborock_future import RoborockFuture
 from .roborock_message import (
@@ -89,20 +87,18 @@ async def validate_connection(self) -> None:
             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)
@@ -115,7 +111,7 @@ def _async_response(
             )
             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))
 
     async def send_message(self, roborock_message: RoborockMessage):
         raise NotImplementedError

roborock/cloud_api.py

@@ -1,11 +1,10 @@
 from __future__ import annotations
 
-import asyncio
 import base64
 import logging
 import threading
 import uuid
-from asyncio import Lock, Task
+from asyncio import Lock
 from typing import Any
 from urllib.parse import urlparse
 
@@ -62,7 +61,7 @@ def on_connect(self, *args, **kwargs):
             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}"
@@ -71,11 +70,11 @@ def on_connect(self, *args, **kwargs):
             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):
         client, __, msg = args
@@ -94,7 +93,7 @@ def on_disconnect(self, *args, **kwargs):
                 self.update_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)
 
@@ -112,53 +111,53 @@ def sync_start_loop(self) -> None:
             self._logger.info("Starting mqtt loop")
             super().loop_start()
 
-    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))
+        disconnected_future = self._async_response(DISCONNECT_REQUEST_ID)
         rc = super().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
+            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))
+        connected_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
+        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:

roborock/roborock_future.py

@@ -14,12 +14,19 @@ def __init__(self, protocol: int):
         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

@@ -378,20 +378,17 @@ def on_message_received(self, messages: list[RoborockMessage]) -> None:
                             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 +448,13 @@ def on_message_received(self, messages: list[RoborockMessage]) -> None:
                         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,8 +1,7 @@
-import asyncio
-
 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 .roborock_client_v1 import COMMANDS_SECURED, RoborockClientV1
@@ -53,16 +52,21 @@ async def send_message(self, roborock_message: RoborockMessage):
         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,4 +1,3 @@
-import asyncio
 import base64
 
 import paho.mqtt.client as mqtt
@@ -10,7 +9,7 @@
 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,
@@ -49,16 +48,21 @@ async def send_message(self, roborock_message: RoborockMessage):
         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

@@ -135,7 +135,7 @@ def on_message_received(self, messages: list[RoborockMessage]) -> None:
                         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)
 
     async def update_values(
         self, dyad_data_protocols: list[RoborockDyadDataProtocol | RoborockZeoProtocol]

roborock/version_a01_apis/roborock_mqtt_client_a01.py

@@ -43,7 +43,7 @@ async def send_message(self, roborock_message: RoborockMessage):
         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 +54,7 @@ async def send_message(self, roborock_message: RoborockMessage):
                     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(

tests/test_api.py

@@ -51,8 +51,7 @@ async def test_can_create_mqtt_roborock():
 async def test_sync_connect(mqtt_client):
     with patch("paho.mqtt.client.Client.connect", return_value=mqtt.MQTT_ERR_SUCCESS):
         with patch("paho.mqtt.client.Client.loop_start", return_value=mqtt.MQTT_ERR_SUCCESS):
-            connecting, connected_future = mqtt_client.sync_connect()
-            assert connecting is True
+            connected_future = mqtt_client.sync_connect()
             assert connected_future is not None
 
             connected_future.cancel()

tests/test_queue.py

@@ -2,6 +2,7 @@
 
 import pytest
 
+from roborock.exceptions import VacuumError
 from roborock.roborock_future import RoborockFuture
 
 
@@ -10,10 +11,18 @@ def test_can_create():
 
 
 @pytest.mark.asyncio
-async def test_put():
+async def test_set_result():
     rq = RoborockFuture(1)
-    rq.resolve(("test", None))
-    assert await rq.async_get(1) == ("test", None)
+    rq.set_result("test")
+    assert await rq.async_get(1) == "test"
+
+
+@pytest.mark.asyncio
+async def test_set_exception():
+    rq = RoborockFuture(1)
+    rq.set_exception(VacuumError("test"))
+    with pytest.raises(VacuumError):
+        assert await rq.async_get(1)
 
 
 @pytest.mark.asyncio