feat: adding status and consumable listeners (#83)

humbertogontijo · web-flow · commit ebdbc907f1f1 · 2023-06-26T16:19:06.000-03:00
* feat: adding status and consumable listeners

* fix: api tests

* chore: linting
diff --git a/roborock/api.py b/roborock/api.py
@@ -55,7 +55,13 @@
 )
 from .protocol import Utils
 from .roborock_future import RoborockFuture
-from .roborock_message import RoborockDataProtocol, RoborockMessage, RoborockMessageProtocol
+from .roborock_message import (
+    ROBOROCK_DATA_CONSUMABLE_PROTOCOL,
+    ROBOROCK_DATA_STATUS_PROTOCOL,
+    RoborockDataProtocol,
+    RoborockMessage,
+    RoborockMessageProtocol,
+)
 from .roborock_typing import DeviceProp, DockSummary, RoborockCommand
 from .util import RepeatableTask, get_running_loop_or_create_one, unpack_list
 
@@ -124,18 +130,23 @@ def stop(self):
         self.task.cancel()
 
     async def update_value(self, params):
+        if self.attribute.set_command is None:
+            raise RoborockException(f"{self.attribute.attribute} have no set command")
         response = await self.api._send_command(self.attribute.set_command, params)
         await self._async_value()
         return response
 
     async def close_value(self):
         if self.attribute.close_command is None:
-            raise RoborockException(f"{self.attribute.attribute} is not closeable")
+            raise RoborockException(f"{self.attribute.attribute} have no close command")
         response = await self.api._send_command(self.attribute.close_command)
         await self._async_value()
         return response
 
 
+device_cache: dict[str, dict[CacheableAttribute, AttributeCache]] = {}
+
+
 class RoborockClient:
     def __init__(self, endpoint: str, device_info: DeviceData) -> None:
         self.event_loop = get_running_loop_or_create_one()
@@ -147,9 +158,15 @@ def __init__(self, endpoint: str, device_info: DeviceData) -> None:
         self._last_disconnection = self.time_func()
         self.keep_alive = KEEPALIVE
         self._diagnostic_data: dict[str, dict[str, Any]] = {}
-        self.cache: dict[CacheableAttribute, AttributeCache] = {
-            cacheable_attribute: AttributeCache(attr, self) for cacheable_attribute, attr in create_cache_map().items()
-        }
+        cache = device_cache.get(device_info.device.duid)
+        if not cache:
+            cache = {
+                cacheable_attribute: AttributeCache(attr, self)
+                for cacheable_attribute, attr in create_cache_map().items()
+            }
+            device_cache[device_info.device.duid] = cache
+        self.cache = cache
+        self._listeners: list[Callable[[CacheableAttribute, RoborockBase], None]] = []
 
     def __del__(self) -> None:
         self.release()
@@ -212,6 +229,30 @@ def on_message_received(self, messages: list[RoborockMessage]) -> None:
                                     if isinstance(result, list) and len(result) == 1:
                                         result = result[0]
                                     queue.resolve((result, None))
+                        else:
+                            try:
+                                data_protocol = RoborockDataProtocol(int(data_point_number))
+                                _LOGGER.debug(f"Got device update for {data_protocol.name}: {data_point}")
+                                if data_protocol in ROBOROCK_DATA_STATUS_PROTOCOL:
+                                    _cls: Type[Status] = ModelStatus.get(
+                                        self.device_info.model, S7MaxVStatus
+                                    )  # Default to S7 MAXV if we don't have the data
+                                    value = self.cache[CacheableAttribute.status].value
+                                    value[data_protocol.name] = data_point
+                                    status = _cls.from_dict(value)
+                                    for listener in self._listeners:
+                                        listener(CacheableAttribute.status, status)
+                                elif data_protocol in ROBOROCK_DATA_CONSUMABLE_PROTOCOL:
+                                    value = self.cache[CacheableAttribute.consumable].value
+                                    value[data_protocol.name] = data_point
+                                    consumable = Consumable.from_dict(value)
+                                    for listener in self._listeners:
+                                        listener(CacheableAttribute.consumable, consumable)
+                                return
+                            except ValueError:
+                                pass
+                            dps = {data_point_number: data_point}
+                            _LOGGER.debug(f"Got unknown data point {dps}")
                 elif data.payload and protocol == RoborockMessageProtocol.MAP_RESPONSE:
                     payload = data.payload[0:24]
                     [endpoint, _, request_id, _] = struct.unpack("<8s8sH6s", payload)
@@ -223,8 +264,6 @@ def on_message_received(self, messages: list[RoborockMessage]) -> None:
                             if isinstance(decompressed, list):
                                 decompressed = decompressed[0]
                             queue.resolve((decompressed, None))
-                elif data.payload and protocol in RoborockDataProtocol:
-                    _LOGGER.debug(f"Got device update for {RoborockDataProtocol(protocol).name}: {data.payload!r}")
                 else:
                     queue = self._waiting_queue.get(data.seq)
                     if queue:
@@ -329,13 +368,15 @@ async def get_status(self) -> Status | None:
         _cls: Type[Status] = ModelStatus.get(
             self.device_info.model, S7MaxVStatus
         )  # Default to S7 MAXV if we don't have the data
-        return await self.send_command(RoborockCommand.GET_STATUS, return_type=_cls)
+        return _cls.from_dict(await self.cache[CacheableAttribute.status].async_value())
 
     async def get_dnd_timer(self) -> DnDTimer | None:
-        return await self.send_command(RoborockCommand.GET_DND_TIMER, return_type=DnDTimer)
+        return DnDTimer.from_dict(await self.cache[CacheableAttribute.dnd_timer].async_value())
 
     async def get_valley_electricity_timer(self) -> ValleyElectricityTimer | None:
-        return await self.send_command(RoborockCommand.GET_VALLEY_ELECTRICITY_TIMER, return_type=ValleyElectricityTimer)
+        return ValleyElectricityTimer.from_dict(
+            await self.cache[CacheableAttribute.valley_electricity_timer].async_value()
+        )
 
     async def get_clean_summary(self) -> CleanSummary | None:
         clean_summary: dict | list | int = await self.send_command(RoborockCommand.GET_CLEAN_SUMMARY)
@@ -357,16 +398,16 @@ async def get_clean_record(self, record_id: int) -> CleanRecord | None:
         return await self.send_command(RoborockCommand.GET_CLEAN_RECORD, [record_id], return_type=CleanRecord)
 
     async def get_consumable(self) -> Consumable | None:
-        return await self.send_command(RoborockCommand.GET_CONSUMABLE, return_type=Consumable)
+        return Consumable.from_dict(await self.cache[CacheableAttribute.consumable].async_value())
 
     async def get_wash_towel_mode(self) -> WashTowelMode | None:
-        return await self.send_command(RoborockCommand.GET_WASH_TOWEL_MODE, return_type=WashTowelMode)
+        return WashTowelMode.from_dict(await self.cache[CacheableAttribute.wash_towel_mode].async_value())
 
     async def get_dust_collection_mode(self) -> DustCollectionMode | None:
-        return await self.send_command(RoborockCommand.GET_DUST_COLLECTION_MODE, return_type=DustCollectionMode)
+        return DustCollectionMode.from_dict(await self.cache[CacheableAttribute.dust_collection_mode].async_value())
 
     async def get_smart_wash_params(self) -> SmartWashParams | None:
-        return await self.send_command(RoborockCommand.GET_SMART_WASH_PARAMS, return_type=SmartWashParams)
+        return SmartWashParams.from_dict(await self.cache[CacheableAttribute.smart_wash_params].async_value())
 
     async def get_dock_summary(self, dock_type: RoborockDockTypeCode) -> DockSummary | None:
         """Gets the status summary from the dock with the methods available for a given dock.
@@ -432,15 +473,18 @@ async def get_room_mapping(self) -> list[RoomMapping] | None:
 
     async def get_child_lock_status(self) -> ChildLockStatus | None:
         """Gets current child lock status."""
-        return await self.send_command(RoborockCommand.GET_CHILD_LOCK_STATUS, return_type=ChildLockStatus)
+        return ChildLockStatus.from_dict(await self.cache[CacheableAttribute.child_lock_status].async_value())
 
     async def get_flow_led_status(self) -> FlowLedStatus | None:
         """Gets current flow led status."""
-        return await self.send_command(RoborockCommand.GET_FLOW_LED_STATUS, return_type=FlowLedStatus)
+        return FlowLedStatus.from_dict(await self.cache[CacheableAttribute.flow_led_status].async_value())
 
     async def get_sound_volume(self) -> int | None:
         """Gets current volume level."""
-        return await self.send_command(RoborockCommand.GET_SOUND_VOLUME)
+        return await self.cache[CacheableAttribute.sound_volume].async_value()
+
+    def add_listener(self, listener: Callable):
+        self._listeners.append(listener)
 
 
 class RoborockApiClient:
diff --git a/roborock/command_cache.py b/roborock/command_cache.py
@@ -11,6 +11,8 @@
 
 
 class CacheableAttribute(str, Enum):
+    status = "status"
+    consumable = "consumable"
     sound_volume = "sound_volume"
     camera_status = "camera_status"
     carpet_clean_mode = "carpet_clean_mode"
@@ -36,12 +38,20 @@ class CacheableAttribute(str, Enum):
 class RoborockAttribute:
     attribute: str
     get_command: RoborockCommand
-    set_command: RoborockCommand
+    set_command: Optional[RoborockCommand] = None
     close_command: Optional[RoborockCommand] = None
 
 
 def create_cache_map():
     cache_map: Mapping[CacheableAttribute, RoborockAttribute] = {
+        CacheableAttribute.status: RoborockAttribute(
+            attribute="status",
+            get_command=RoborockCommand.GET_STATUS,
+        ),
+        CacheableAttribute.consumable: RoborockAttribute(
+            attribute="consumable",
+            get_command=RoborockCommand.GET_CONSUMABLE,
+        ),
         CacheableAttribute.sound_volume: RoborockAttribute(
             attribute="sound_volume",
             get_command=RoborockCommand.GET_SOUND_VOLUME,
diff --git a/roborock/roborock_message.py b/roborock/roborock_message.py
@@ -28,9 +28,9 @@ class RoborockDataProtocol(RoborockEnum):
     BATTERY = 122
     FAN_POWER = 123
     WATER_BOX_MODE = 124
-    MAIN_BRUSH_LIFE = 125
-    SIDE_BRUSH_LIFE = 126
-    FILTER_LIFE = 127
+    MAIN_BRUSH_WORK_TIME = 125
+    SIDE_BRUSH_WORK_TIME = 126
+    FILTER_WORK_TIME = 127
     ADDITIONAL_PROPS = 128
     TASK_COMPLETE = 130
     TASK_CANCEL_LOW_POWER = 131
@@ -39,6 +39,22 @@ class RoborockDataProtocol(RoborockEnum):
     DRYING_STATUS = 134
 
 
+ROBOROCK_DATA_STATUS_PROTOCOL = [
+    RoborockDataProtocol.ERROR_CODE,
+    RoborockDataProtocol.STATE,
+    RoborockDataProtocol.BATTERY,
+    RoborockDataProtocol.FAN_POWER,
+    RoborockDataProtocol.WATER_BOX_MODE,
+    RoborockDataProtocol.CHARGE_STATUS,
+]
+
+ROBOROCK_DATA_CONSUMABLE_PROTOCOL = [
+    RoborockDataProtocol.MAIN_BRUSH_WORK_TIME,
+    RoborockDataProtocol.SIDE_BRUSH_WORK_TIME,
+    RoborockDataProtocol.FILTER_WORK_TIME,
+]
+
+
 @dataclass
 class RoborockMessage:
     protocol: RoborockMessageProtocol
diff --git a/tests/test_api.py b/tests/test_api.py
@@ -12,7 +12,7 @@
 )
 from roborock.api import PreparedRequest, RoborockApiClient
 from roborock.cloud_api import RoborockMqttClient
-from roborock.containers import DeviceData, DustCollectionMode, S7MaxVStatus, SmartWashParams, WashTowelMode
+from roborock.containers import DeviceData, S7MaxVStatus
 from tests.mock_data import BASE_URL_REQUEST, GET_CODE_RESPONSE, HOME_DATA_RAW, STATUS, USER_DATA
 
 
@@ -77,8 +77,8 @@ async def test_get_dust_collection_mode():
     home_data = HomeData.from_dict(HOME_DATA_RAW)
     device_info = DeviceData(device=home_data.devices[0], model=home_data.products[0].model)
     rmc = RoborockMqttClient(UserData.from_dict(USER_DATA), device_info)
-    with patch("roborock.cloud_api.RoborockMqttClient.send_command") as command:
-        command.return_value = DustCollectionMode.from_dict({"mode": 1})
+    with patch("roborock.api.AttributeCache.async_value") as command:
+        command.return_value = {"mode": 1}
         dust = await rmc.get_dust_collection_mode()
         assert dust is not None
         assert dust.mode == RoborockDockDustCollectionModeCode.light
@@ -89,8 +89,8 @@ async def test_get_mop_wash_mode():
     home_data = HomeData.from_dict(HOME_DATA_RAW)
     device_info = DeviceData(device=home_data.devices[0], model=home_data.products[0].model)
     rmc = RoborockMqttClient(UserData.from_dict(USER_DATA), device_info)
-    with patch("roborock.cloud_api.RoborockMqttClient.send_command") as command:
-        command.return_value = SmartWashParams.from_dict({"smart_wash": 0, "wash_interval": 1500})
+    with patch("roborock.api.AttributeCache.async_value") as command:
+        command.return_value = {"smart_wash": 0, "wash_interval": 1500}
         mop_wash = await rmc.get_smart_wash_params()
         assert mop_wash is not None
         assert mop_wash.smart_wash == 0
@@ -102,8 +102,8 @@ async def test_get_washing_mode():
     home_data = HomeData.from_dict(HOME_DATA_RAW)
     device_info = DeviceData(device=home_data.devices[0], model=home_data.products[0].model)
     rmc = RoborockMqttClient(UserData.from_dict(USER_DATA), device_info)
-    with patch("roborock.cloud_api.RoborockMqttClient.send_command") as command:
-        command.return_value = WashTowelMode.from_dict({"wash_mode": 2})
+    with patch("roborock.api.AttributeCache.async_value") as command:
+        command.return_value = {"wash_mode": 2}
         washing_mode = await rmc.get_wash_towel_mode()
         assert washing_mode is not None
         assert washing_mode.wash_mode == RoborockDockWashTowelModeCode.deep
@@ -116,8 +116,10 @@ async def test_get_prop():
     device_info = DeviceData(device=home_data.devices[0], model=home_data.products[0].model)
     rmc = RoborockMqttClient(UserData.from_dict(USER_DATA), device_info)
     with patch("roborock.cloud_api.RoborockMqttClient.get_status") as get_status, patch(
-        "roborock.cloud_api.RoborockMqttClient.send_command"
-    ), patch("roborock.cloud_api.RoborockMqttClient.get_dust_collection_mode"):
+        "roborock.api.RoborockClient.send_command"
+    ), patch("roborock.api.AttributeCache.async_value"), patch(
+        "roborock.cloud_api.RoborockMqttClient.get_dust_collection_mode"
+    ):
         status = S7MaxVStatus.from_dict(STATUS)
         status.dock_type = RoborockDockTypeCode.auto_empty_dock_pure
         get_status.return_value = status
