feat: improving cache and refactoring (#82)

Author: humbertogontijo

roborock/api.py

@@ -18,7 +18,7 @@
 import aiohttp
 
 from .code_mappings import RoborockDockTypeCode
-from .command_cache import AttributeCache, CacheableAttribute, CommandType, parse_method
+from .command_cache import CacheableAttribute, CommandType, RoborockAttribute, create_cache_map, parse_method
 from .containers import (
     ChildLockStatus,
     CleanRecord,
@@ -55,9 +55,9 @@
 )
 from .protocol import Utils
 from .roborock_future import RoborockFuture
-from .roborock_message import RoborockDataProtocol, RoborockMessage
+from .roborock_message import RoborockDataProtocol, RoborockMessage, RoborockMessageProtocol
 from .roborock_typing import DeviceProp, DockSummary, RoborockCommand
-from .util import unpack_list
+from .util import RepeatableTask, get_running_loop_or_create_one, unpack_list
 
 _LOGGER = logging.getLogger(__name__)
 KEEPALIVE = 60
@@ -94,8 +94,42 @@ async def request(self, method: str, url: str, params=None, data=None, headers=N
                 return await resp.json()
 
 
+EVICT_TIME = 60
+
+
+class AttributeCache:
+    def __init__(self, attribute: RoborockAttribute, api: RoborockClient):
+        self.attribute = attribute
+        self.api = api
+        self.attribute = attribute
+        self.task = RepeatableTask(self.api.event_loop, self._async_value, EVICT_TIME)
+        self._value: Any = None
+
+    @property
+    def value(self):
+        return self._value
+
+    async def _async_value(self):
+        self._value = await self.api._send_command(self.attribute.get_command)
+        return self._value
+
+    async def async_value(self):
+        if self._value is None:
+            return await self.task.reset()
+        return self._value
+
+    def stop(self):
+        self.task.cancel()
+
+    async def update_value(self, params):
+        response = await self.api._send_command(self.attribute.set_command, params)
+        await self._async_value()
+        return response
+
+
 class RoborockClient:
     def __init__(self, endpoint: str, device_info: DeviceData) -> None:
+        self.event_loop = get_running_loop_or_create_one()
         self.device_info = device_info
         self._endpoint = endpoint
         self._nonce = secrets.token_bytes(16)
@@ -105,11 +139,15 @@ def __init__(self, endpoint: str, device_info: DeviceData) -> None:
         self.keep_alive = KEEPALIVE
         self._diagnostic_data: dict[str, dict[str, Any]] = {}
         self.cache: dict[CacheableAttribute, AttributeCache] = {
-            cacheable_attribute: AttributeCache(cacheable_attribute) for cacheable_attribute in CacheableAttribute
+            cacheable_attribute: AttributeCache(attr, self) for cacheable_attribute, attr in create_cache_map().items()
         }
 
     def __del__(self) -> None:
+        self.release()
+
+    def release(self):
         self.sync_disconnect()
+        [item.stop() for item in self.cache.values()]
 
     @property
     def diagnostic_data(self) -> dict:
@@ -138,7 +176,10 @@ def on_message_received(self, messages: list[RoborockMessage]) -> None:
             self._last_device_msg_in = self.time_func()
             for data in messages:
                 protocol = data.protocol
-                if data.payload and (protocol == 102 or protocol == 4):
+                if data.payload and protocol in [
+                    RoborockMessageProtocol.RPC_RESPONSE,
+                    RoborockMessageProtocol.GENERAL_REQUEST,
+                ]:
                     payload = json.loads(data.payload.decode())
                     for data_point_number, data_point in payload.get("dps").items():
                         if data_point_number == "102":
@@ -162,7 +203,7 @@ def on_message_received(self, messages: list[RoborockMessage]) -> None:
                                     if isinstance(result, list) and len(result) == 1:
                                         result = result[0]
                                     queue.resolve((result, None))
-                elif data.payload and protocol == 301:
+                elif data.payload and protocol == RoborockMessageProtocol.MAP_RESPONSE:
                     payload = data.payload[0:24]
                     [endpoint, _, request_id, _] = struct.unpack("<8s8sH6s", payload)
                     if endpoint.decode().startswith(self._endpoint):
@@ -242,6 +283,9 @@ def _get_payload(
         )
         return request_id, timestamp, payload
 
+    async def send_message(self, roborock_message: RoborockMessage):
+        raise NotImplementedError
+
     async def _send_command(
         self,
         method: RoborockCommand,
@@ -257,20 +301,17 @@ async def send_command(
         return_type: Optional[Type[RT]] = None,
     ) -> RT:
         parsed_method = parse_method(method)
-        if parsed_method is not None and parsed_method.type == CommandType.GET:
-            cache = self.cache[parsed_method.attribute].value
-            if cache is not None:
-                if return_type:
-                    return return_type.from_dict(cache)
-                return cache
-
-        response = await self._send_command(method, params)
-
+        cache = None
         if parsed_method is not None:
-            if parsed_method.type == CommandType.SET:
-                self.cache[parsed_method.attribute].evict()
-            elif parsed_method.type == CommandType.GET:
-                self.cache[parsed_method.attribute].load(response)
+            cache = self.cache[parsed_method.attribute]
+        response: Any = None
+        if cache is not None:
+            if parsed_method.type == CommandType.GET:
+                response = await cache.async_value()
+            elif parsed_method.type == CommandType.SET:
+                response = await cache.update_value(params)
+        else:
+            response = await self._send_command(method, params)
         if return_type:
             return return_type.from_dict(response)
         return response

roborock/cloud_api.py

@@ -15,7 +15,7 @@
 from .exceptions import CommandVacuumError, RoborockException, VacuumError
 from .protocol import MessageParser, Utils
 from .roborock_future import RoborockFuture
-from .roborock_message import RoborockMessage
+from .roborock_message import RoborockMessage, RoborockMessageProtocol
 from .roborock_typing import RoborockCommand
 
 _LOGGER = logging.getLogger(__name__)
@@ -149,33 +149,44 @@ def _send_msg_raw(self, msg: bytes) -> None:
         if info.rc != mqtt.MQTT_ERR_SUCCESS:
             raise RoborockException(f"Failed to publish ({mqtt.error_string(info.rc)})")
 
-    async def _send_command(
-        self,
-        method: RoborockCommand,
-        params: Optional[list | dict] = None,
-    ):
+    async def send_message(self, roborock_message: RoborockMessage):
         await self.validate_connection()
-        request_id, timestamp, payload = super()._get_payload(method, params, True)
-        _LOGGER.debug(f"id={request_id} Requesting method {method} with {params}")
-        request_protocol = 101
-        response_protocol = 301 if method in COMMANDS_SECURED else 102
-        roborock_message = RoborockMessage(timestamp=timestamp, protocol=request_protocol, payload=payload)
+        method = roborock_message.get_method()
+        params = roborock_message.get_params()
+        request_id = roborock_message.get_request_id()
+        if request_id is None:
+            raise RoborockException(f"Failed build message {roborock_message}")
+        response_protocol = (
+            RoborockMessageProtocol.MAP_RESPONSE if method in COMMANDS_SECURED else RoborockMessageProtocol.RPC_RESPONSE
+        )
+
         local_key = self.device_info.device.local_key
         msg = MessageParser.build(roborock_message, local_key, False)
+        _LOGGER.debug(f"id={request_id} Requesting method {method} with {params}")
         self._send_msg_raw(msg)
         (response, err) = await self._async_response(request_id, response_protocol)
         self._diagnostic_data[method if method is not None else "unknown"] = {
-            "params": params,
+            "params": roborock_message.get_params(),
             "response": response,
             "error": err,
         }
         if err:
             raise CommandVacuumError(method, err) from err
-        if response_protocol == 301:
+        if response_protocol == RoborockMessageProtocol.MAP_RESPONSE:
             _LOGGER.debug(f"id={request_id} Response from {method}: {len(response)} bytes")
         else:
             _LOGGER.debug(f"id={request_id} Response from {method}: {response}")
         return response
 
+    async def _send_command(
+        self,
+        method: RoborockCommand,
+        params: Optional[list | dict] = None,
+    ):
+        request_id, timestamp, payload = super()._get_payload(method, params, True)
+        request_protocol = RoborockMessageProtocol.RPC_REQUEST
+        roborock_message = RoborockMessage(timestamp=timestamp, protocol=request_protocol, payload=payload)
+        return await self.send_message(roborock_message)
+
     async def get_map_v1(self):
         return await self.send_command(RoborockCommand.GET_MAP_V1)

roborock/command_cache.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
-import time
 from dataclasses import dataclass
 from enum import Enum
+from typing import Mapping
+
+from roborock import RoborockCommand
 
 GET_PREFIX = "get_"
 SET_PREFIX = ("set_", "change_", "close_")
@@ -30,29 +32,108 @@ class CacheableAttribute(str, Enum):
     wash_towel_mode = "wash_towel_mode"
 
 
-EVICT_TIME = 60
-
-
-class AttributeCache:
-    _value: dict | None = None
-    last_update: float | None = None
-
-    def __init__(self, attribute: str):
-        self.attribute = attribute
-
-    @property
-    def value(self):
-        if self.last_update is None or time.monotonic() - self.last_update > EVICT_TIME:
-            self._value = None
-        return self._value
-
-    def load(self, value: dict):
-        self._value = value
-        self.last_update = time.monotonic()
-        return self._value
-
-    def evict(self):
-        self._value = None
+@dataclass
+class RoborockAttribute:
+    attribute: str
+    get_command: RoborockCommand
+    set_command: RoborockCommand
+
+
+def create_cache_map():
+    cache_map: Mapping[CacheableAttribute, RoborockAttribute] = {
+        CacheableAttribute.sound_volume: RoborockAttribute(
+            attribute="sound_volume",
+            get_command=RoborockCommand.GET_SOUND_VOLUME,
+            set_command=RoborockCommand.CHANGE_SOUND_VOLUME,
+        ),
+        CacheableAttribute.camera_status: RoborockAttribute(
+            attribute="camera_status",
+            get_command=RoborockCommand.GET_CAMERA_STATUS,
+            set_command=RoborockCommand.SET_CAMERA_STATUS,
+        ),
+        CacheableAttribute.carpet_clean_mode: RoborockAttribute(
+            attribute="carpet_clean_mode",
+            get_command=RoborockCommand.GET_CARPET_CLEAN_MODE,
+            set_command=RoborockCommand.SET_CARPET_CLEAN_MODE,
+        ),
+        CacheableAttribute.carpet_mode: RoborockAttribute(
+            attribute="carpet_mode",
+            get_command=RoborockCommand.GET_CARPET_MODE,
+            set_command=RoborockCommand.SET_CARPET_MODE,
+        ),
+        CacheableAttribute.child_lock_status: RoborockAttribute(
+            attribute="child_lock_status",
+            get_command=RoborockCommand.GET_CHILD_LOCK_STATUS,
+            set_command=RoborockCommand.SET_CHILD_LOCK_STATUS,
+        ),
+        CacheableAttribute.collision_avoid_status: RoborockAttribute(
+            attribute="collision_avoid_status",
+            get_command=RoborockCommand.GET_COLLISION_AVOID_STATUS,
+            set_command=RoborockCommand.SET_COLLISION_AVOID_STATUS,
+        ),
+        CacheableAttribute.customize_clean_mode: RoborockAttribute(
+            attribute="customize_clean_mode",
+            get_command=RoborockCommand.GET_CUSTOMIZE_CLEAN_MODE,
+            set_command=RoborockCommand.SET_CUSTOMIZE_CLEAN_MODE,
+        ),
+        CacheableAttribute.custom_mode: RoborockAttribute(
+            attribute="custom_mode",
+            get_command=RoborockCommand.GET_CUSTOM_MODE,
+            set_command=RoborockCommand.SET_CUSTOM_MODE,
+        ),
+        CacheableAttribute.dnd_timer: RoborockAttribute(
+            attribute="dnd_timer", get_command=RoborockCommand.GET_DND_TIMER, set_command=RoborockCommand.SET_DND_TIMER
+        ),
+        CacheableAttribute.dust_collection_mode: RoborockAttribute(
+            attribute="dust_collection_mode",
+            get_command=RoborockCommand.GET_DUST_COLLECTION_MODE,
+            set_command=RoborockCommand.SET_DUST_COLLECTION_MODE,
+        ),
+        CacheableAttribute.flow_led_status: RoborockAttribute(
+            attribute="flow_led_status",
+            get_command=RoborockCommand.GET_FLOW_LED_STATUS,
+            set_command=RoborockCommand.SET_FLOW_LED_STATUS,
+        ),
+        CacheableAttribute.identify_furniture_status: RoborockAttribute(
+            attribute="identify_furniture_status",
+            get_command=RoborockCommand.GET_IDENTIFY_FURNITURE_STATUS,
+            set_command=RoborockCommand.SET_IDENTIFY_FURNITURE_STATUS,
+        ),
+        CacheableAttribute.identify_ground_material_status: RoborockAttribute(
+            attribute="identify_ground_material_status",
+            get_command=RoborockCommand.GET_IDENTIFY_GROUND_MATERIAL_STATUS,
+            set_command=RoborockCommand.SET_IDENTIFY_GROUND_MATERIAL_STATUS,
+        ),
+        CacheableAttribute.led_status: RoborockAttribute(
+            attribute="led_status",
+            get_command=RoborockCommand.GET_LED_STATUS,
+            set_command=RoborockCommand.SET_LED_STATUS,
+        ),
+        CacheableAttribute.server_timer: RoborockAttribute(
+            attribute="server_timer",
+            get_command=RoborockCommand.GET_SERVER_TIMER,
+            set_command=RoborockCommand.SET_SERVER_TIMER,
+        ),
+        CacheableAttribute.smart_wash_params: RoborockAttribute(
+            attribute="smart_wash_params",
+            get_command=RoborockCommand.GET_SMART_WASH_PARAMS,
+            set_command=RoborockCommand.SET_SMART_WASH_PARAMS,
+        ),
+        CacheableAttribute.timezone: RoborockAttribute(
+            attribute="timezone", get_command=RoborockCommand.GET_TIMEZONE, set_command=RoborockCommand.SET_TIMEZONE
+        ),
+        CacheableAttribute.valley_electricity_timer: RoborockAttribute(
+            attribute="valley_electricity_timer",
+            get_command=RoborockCommand.GET_VALLEY_ELECTRICITY_TIMER,
+            set_command=RoborockCommand.SET_VALLEY_ELECTRICITY_TIMER,
+        ),
+        CacheableAttribute.wash_towel_mode: RoborockAttribute(
+            attribute="wash_towel_mode",
+            get_command=RoborockCommand.GET_WASH_TOWEL_MODE,
+            set_command=RoborockCommand.SET_WASH_TOWEL_MODE,
+        ),
+    }
+    return cache_map
 
 
 class CommandType(Enum):
@@ -78,9 +159,7 @@ def parse_method(method: str):
             for prefix in SET_PREFIX:
                 attribute = attribute.removeprefix(prefix)
             command_type = CommandType.SET
-        try:
-            cacheable_attribute = CacheableAttribute(attribute)
-            return ParserCommand(type=command_type, attribute=cacheable_attribute)
-        except ValueError:
-            pass
+        cacheable_attribute = next((attr for attr in CacheableAttribute if attr == attribute), None)
+        if cacheable_attribute:
+            return ParserCommand(type=command_type, attribute=CacheableAttribute(cacheable_attribute))
     return None

roborock/exceptions.py

@@ -1,4 +1,5 @@
 """Roborock exceptions."""
+from __future__ import annotations
 
 
 class RoborockException(Exception):
@@ -24,8 +25,8 @@ class VacuumError(RoborockException):
 class CommandVacuumError(RoborockException):
     """Class for command vacuum errors."""
 
-    def __init__(self, command: str, vacuum_error: VacuumError):
-        self.message = f"{command}: {str(vacuum_error)}"
+    def __init__(self, command: str | None, vacuum_error: VacuumError):
+        self.message = f"{command or 'unknown'}: {str(vacuum_error)}"
         super().__init__(self.message)