Merge pull request #10 from Lash-L/docks

Add Docking Information

Author: humbertogontijo

roborock/api.py

@@ -1,4 +1,5 @@
 """The Roborock api."""
+
 from __future__ import annotations
 
 import asyncio
@@ -23,9 +24,16 @@
 import paho.mqtt.client as mqtt
 from Crypto.Cipher import AES
 from Crypto.Util.Padding import pad, unpad
-from roborock.code_mappings import STATE_CODE_TO_STATUS
 
-from roborock.containers import (
+from roborock.exceptions import (
+    RoborockException,
+    CommandVacuumError,
+    VacuumError,
+    RoborockTimeout,
+)
+from .code_mappings import STATE_CODE_TO_STATUS, WASH_MODE_MAP, DUST_COLLECTION_MAP, RoborockDockType, \
+    RoborockDockDustCollectionType, RoborockDockWashingModeType
+from .containers import (
     UserData,
     HomeDataDevice,
     Status,
@@ -35,20 +43,17 @@
     CleanRecord,
     HomeData,
     MultiMapsList,
+    SmartWashParameters,
+
 )
-from roborock.exceptions import (
-    RoborockException,
-    CommandVacuumError,
-    VacuumError,
-    RoborockTimeout,
-)
-from roborock.roborock_queue import RoborockQueue
-from roborock.typing import (
+from .roborock_queue import RoborockQueue
+from .typing import (
     RoborockDeviceInfo,
     RoborockDeviceProp,
     RoborockCommand,
+    RoborockDockSummary,
 )
-from roborock.util import run_in_executor
+from .util import run_in_executor
 
 _LOGGER = logging.getLogger(__name__)
 QUEUE_TIMEOUT = 4
@@ -78,17 +83,17 @@ def __init__(self, base_url: str, base_headers: dict = None) -> None:
         self.base_headers = base_headers or {}
 
     async def request(
-            self, method: str, url: str, params=None, data=None, headers=None
+        self, method: str, url: str, params=None, data=None, headers=None
     ) -> dict | list:
         _url = "/".join(s.strip("/") for s in [self.base_url, url])
         _headers = {**self.base_headers, **(headers or {})}
         async with aiohttp.ClientSession() as session:
             async with session.request(
-                    method,
-                    _url,
-                    params=params,
-                    data=data,
-                    headers=_headers,
+                method,
+                _url,
+                params=params,
+                data=data,
+                headers=_headers,
             ) as resp:
                 return await resp.json()
 
@@ -369,7 +374,7 @@ def _send_msg_raw(self, device_id, protocol, timestamp, payload) -> None:
             raise RoborockException(f"Failed to publish (rc: {info.rc})")
 
     async def send_command(
-            self, device_id: str, method: RoborockCommand, params: list = None
+        self, device_id: str, method: RoborockCommand, params: list = None
     ):
         await self.validate_connection()
         timestamp = math.floor(time.time())
@@ -439,6 +444,29 @@ async def get_consumable(self, device_id: str) -> Consumable:
         if isinstance(consumable, dict):
             return Consumable(consumable)
 
+    async def get_washing_mode(self, device_id: str) -> RoborockDockWashingModeType:
+        washing_mode = await self.send_command(device_id, RoborockCommand.GET_WASH_TOWEL_MODE)
+        return WASH_MODE_MAP.get(washing_mode)
+
+    async def get_dust_collection_mode(self, device_id: str) -> RoborockDockDustCollectionType:
+        dust_collection = await self.send_command(device_id, RoborockCommand.GET_DUST_COLLECTION_MODE)
+        return DUST_COLLECTION_MAP.get(dust_collection)
+
+    async def get_mop_wash_mode(self, device_id: str) -> SmartWashParameters:
+        mop_wash_mode = await self.send_command(device_id, RoborockCommand.GET_SMART_WASH_PARAMS)
+        if isinstance(mop_wash_mode, dict):
+            return SmartWashParameters(mop_wash_mode)
+
+    async def get_dock_summary(self, device_id: str, dock_type: RoborockDockType) -> RoborockDockSummary:
+        collection_mode = await self.get_dust_collection_mode(device_id)
+        mop_wash = None
+        washing_mode = None
+        if dock_type == RoborockDockType.EMPTY_WASH_FILL_DOCK:
+            [mop_wash, washing_mode] = await asyncio.gather(
+                *[self.get_mop_wash_mode(device_id), self.get_washing_mode(device_id)])
+
+        return RoborockDockSummary(collection_mode, washing_mode, mop_wash)
+
     async def get_prop(self, device_id: str) -> RoborockDeviceProp:
         [status, dnd_timer, clean_summary, consumable] = await asyncio.gather(
             *[
@@ -453,9 +481,12 @@ async def get_prop(self, device_id: str) -> RoborockDeviceProp:
             last_clean_record = await self.get_clean_record(
                 device_id, clean_summary.records[0]
             )
+        dock_summary = None
+        if status.dock_type != RoborockDockType.NO_DOCK:
+            dock_summary = await self.get_dock_summary(device_id, status.dock_type)
         if any([status, dnd_timer, clean_summary, consumable]):
             return RoborockDeviceProp(
-                status, dnd_timer, clean_summary, consumable, last_clean_record
+                status, dnd_timer, clean_summary, consumable, last_clean_record, dock_summary
             )
 
     async def get_multi_maps_list(self, device_id) -> MultiMapsList:

roborock/code_mappings.py

@@ -1,7 +1,31 @@
 from __future__ import annotations
 
+from enum import Enum
 from typing import Any
 
+
+class RoborockDockType(str, Enum):
+    NO_DOCK = "No dock"
+    AUTO_EMPTY = "Roborock auto-empty dock"  # Dust collection
+    EMPTY_WASH_FILL_DOCK = "Roborock empty wash fill dock"
+    AUTO_EMPTY_PURE = "Roborock auto-empty pure dock"
+    UNKNOWN = "Unknown Dock - please submit an issue"
+
+
+class RoborockDockDustCollectionType(str, Enum):
+    SMART = "smart"
+    QUICK = "quick"
+    DAILY = "daily"
+    STRONG = "strong"
+    MAX = "max"
+
+
+class RoborockDockWashingModeType(str, Enum):
+    LIGHT = "light"
+    BALANCED = "balanced"
+    DEEP = "deep"
+
+
 STATE_CODE_TO_STATUS: dict[int | Any, str | Any] = {
     1: "starting",
     2: "charger_disconnected",
@@ -82,3 +106,28 @@
     203: "intense",
     204: "custom",
 }
+
+DOCK_ERROR_TO_TEXT = {
+    0: "ok",
+    38: 'water empty',
+    39: 'waste water tank full',
+}
+
+DOCK_TYPE_MAP = {
+    0: RoborockDockType.NO_DOCK,
+    3: RoborockDockType.EMPTY_WASH_FILL_DOCK,
+}
+
+DUST_COLLECTION_MAP = {
+    0: RoborockDockDustCollectionType.SMART,
+    1: RoborockDockDustCollectionType.QUICK,
+    2: RoborockDockDustCollectionType.DAILY,
+    3: RoborockDockDustCollectionType.STRONG,
+    4: RoborockDockDustCollectionType.MAX,
+}
+
+WASH_MODE_MAP = {
+    0: RoborockDockWashingModeType.LIGHT,
+    1: RoborockDockWashingModeType.BALANCED,
+    2: RoborockDockWashingModeType.DEEP,
+}

roborock/containers.py

@@ -1,7 +1,7 @@
 from enum import Enum
 
-from roborock.code_mappings import STATE_CODE_TO_STATUS, ERROR_CODE_TO_TEXT, FAN_SPEED_CODES, MOP_MODE_CODES, \
-    MOP_INTENSITY_CODES
+from .code_mappings import STATE_CODE_TO_STATUS, ERROR_CODE_TO_TEXT, FAN_SPEED_CODES, MOP_MODE_CODES, \
+    MOP_INTENSITY_CODES, DOCK_ERROR_TO_TEXT, DOCK_TYPE_MAP, RoborockDockType
 
 
 class UserDataRRiotReferenceField(str, Enum):
@@ -220,6 +220,15 @@ class MultiMapListField(str, Enum):
     MAP_INFO = "map_info"
 
 
+class SmartWashField(str, Enum):
+    SMART_WASH = "smart_wash"
+    WASH_INTERVAL = "wash_interval"
+
+
+class WashTowelField(str, Enum):
+    Wash_MODE = "wash_mode"
+
+
 class RoborockBase(dict):
     def __init__(self, data: dict[str, any]) -> None:
         super().__init__()
@@ -761,9 +770,13 @@ def water_shortage_status(self) -> int:
         return self.get(StatusField.WATER_SHORTAGE_STATUS)
 
     @property
-    def dock_type(self) -> int:
+    def dock_type_code(self) -> int:
         return self.get(StatusField.DOCK_TYPE)
 
+    @property
+    def dock_type(self) -> RoborockDockType:
+        return DOCK_TYPE_MAP.get(self.get(StatusField.DOCK_TYPE), RoborockDockType.UNKNOWN)
+
     @property
     def dust_collection_status(self) -> int:
         return self.get(StatusField.DUST_COLLECTION_STATUS)
@@ -797,9 +810,13 @@ def switch_map_mode(self) -> int:
         return self.get(StatusField.SWITCH_MAP_MODE)
 
     @property
-    def dock_error_status(self) -> int:
+    def dock_error_status_code(self) -> int:
         return self.get(StatusField.DOCK_ERROR_STATUS)
 
+    @property
+    def dock_error_status(self) -> str:
+        return DOCK_ERROR_TO_TEXT.get(self.get(StatusField.DOCK_ERROR_STATUS))
+
     @property
     def charge_status(self) -> int:
         return self.get(StatusField.CHARGE_STATUS)
@@ -1014,3 +1031,16 @@ def multi_map_count(self) -> int:
     @property
     def map_info(self) -> list[MultiMapsListMapInfo]:
         return [MultiMapsListMapInfo(map_info) for map_info in self.get(MultiMapListField.MAP_INFO)]
+
+
+class SmartWashParameters(RoborockBase):
+    def __init__(self, data: dict[str, any]) -> None:
+        super().__init__(data)
+
+    @property
+    def smart_wash(self) -> int:
+        return self.get(SmartWashField.SMART_WASH)
+
+    @property
+    def wash_interval(self) -> int:
+        return self.get(SmartWashField.WASH_INTERVAL)

roborock/typing.py

@@ -1,7 +1,8 @@
 from enum import Enum
 
-from roborock.containers import HomeDataDevice, HomeDataProduct, Status, CleanSummary, Consumable, \
-    DNDTimer, CleanRecord
+from .code_mappings import RoborockDockDustCollectionType, RoborockDockWashingModeType
+from .containers import HomeDataDevice, HomeDataProduct, Status, CleanSummary, Consumable, \
+    DNDTimer, CleanRecord, SmartWashParameters
 
 
 class RoborockDevicePropField(str, Enum):
@@ -10,6 +11,14 @@ class RoborockDevicePropField(str, Enum):
     CLEAN_SUMMARY = "clean_summary"
     CONSUMABLE = "consumable"
     LAST_CLEAN_RECORD = "last_clean_record"
+    DOCK_SUMMARY = "dock_summary"
+
+
+class RoborockDockSummaryField(str, Enum):
+    DUST_COLLECTION_MODE = "dust_collection_mode"
+    WASHING_MODE_TYPE = "washing_mode_type"
+    MOP_WASH = "mop_wash"
+
 
 class RoborockCommand(str, Enum):
     GET_MAP_V1 = "get_map_v1",
@@ -38,6 +47,19 @@ class RoborockCommand(str, Enum):
     APP_GOTO_TARGET = "app_goto_target",
     APP_SEGMENT_CLEAN = "app_segment_clean",
     APP_ZONED_CLEAN = "app_zoned_clean",
+    APP_GET_DRYER_SETTING = "app_get_dryer_setting"
+    APP_SET_DRYER_SETTING = "app_set_dryer_setting"
+    APP_START_WASH = "app_start_wash"
+    APP_STOP_WASH = "app_stop_wash"
+    GET_DUST_COLLECTION_MODE = "get_dust_collection_mode"
+    SET_DUST_COLLECTION_MODE = "set_dust_collection_mode"
+    GET_SMART_WASH_PARAMS = "get_smart_wash_params"
+    SET_SMART_WASH_PARAMS = "set_smart_wash_params"
+    GET_WASH_TOWEL_MODE = "get_wash_towel_mode"
+    SET_WASH_TOWEL_MODE = "set_wash_towel_mode"
+    SET_CHILD_LOCK_STATUS = "set_child_lock_status"
+    GET_CHILD_LOCK_STATUS = "get_child_lock_status"
+    START_WASH_THEN_CHARGE = "start_wash_then_charge"
 
 
 class RoborockDeviceInfo:
@@ -46,14 +68,23 @@ def __init__(self, device: HomeDataDevice, product: HomeDataProduct):
         self.product = product
 
 
+class RoborockDockSummary:
+    def __init__(self, dust_collection_mode: RoborockDockDustCollectionType,
+                 washing_mode_type: RoborockDockWashingModeType, mop_wash: SmartWashParameters) -> None:
+        self.dust_collection_mode = dust_collection_mode
+        self.washing_mode_type = washing_mode_type
+        self.mop_wash = mop_wash
+
+
 class RoborockDeviceProp:
     def __init__(self, status: Status, dnd_timer: DNDTimer, clean_summary: CleanSummary, consumable: Consumable,
-                 last_clean_record: CleanRecord):
+                 last_clean_record: CleanRecord, dock_summary: RoborockDockSummary):
         self.status = status
         self.dnd_timer = dnd_timer
         self.clean_summary = clean_summary
         self.consumable = consumable
         self.last_clean_record = last_clean_record
+        self.dock_summary = dock_summary
 
     def update(self, device_prop: 'RoborockDeviceProp'):
         if device_prop.status:
@@ -66,3 +97,5 @@ def update(self, device_prop: 'RoborockDeviceProp'):
             self.consumable = device_prop.consumable
         if device_prop.last_clean_record:
             self.last_clean_record = device_prop.last_clean_record
+        if device_prop.dock_summary:
+            self.dock_summary = device_prop.dock_summary

tests/test_containers.py

@@ -1,4 +1,4 @@
-from roborock import UserData, HomeData, Consumable, Status, DNDTimer, CleanSummary, CleanRecord
+from roborock import UserData, HomeData, Consumable, Status, DNDTimer, CleanSummary, CleanRecord, RoborockDockType
 from .mock_data import USER_DATA, HOME_DATA_RAW, CONSUMABLE, STATUS, DND_TIMER, CLEAN_SUMMARY, CLEAN_RECORD
 
 
@@ -124,7 +124,8 @@ def test_status():
     assert s.home_sec_enable_password == 0
     assert s.adbumper_status == [0, 0, 0]
     assert s.water_shortage_status == 0
-    assert s.dock_type == 3
+    assert s.dock_type_code == 3
+    assert s.dock_type == RoborockDockType.EMPTY_WASH_FILL_DOCK
     assert s.dust_collection_status == 0
     assert s.auto_dust_collection == 1
     assert s.avoid_count == 19
@@ -133,7 +134,8 @@ def test_status():
     assert s.debug_mode == 0
     assert s.collision_avoid_status == 1
     assert s.switch_map_mode == 0
-    assert s.dock_error_status == 0
+    assert s.dock_error_status_code == 0
+    assert s.dock_error_status == "ok"
     assert s.charge_status == 1
     assert s.unsave_map_reason == 0
     assert s.unsave_map_flag == 0