Merge pull request #21 from Lash-L/dock_cleanup

Dock cleanup

Author: humbertogontijo

roborock/api.py

@@ -23,8 +23,8 @@
 from roborock.exceptions import (
     RoborockException, RoborockTimeout, VacuumError,
 )
-from .code_mappings import WASH_MODE_MAP, DUST_COLLECTION_MAP, RoborockDockType, \
-    RoborockDockDustCollectionType, RoborockDockWashingModeType, STATE_CODE_TO_STATUS
+from .code_mappings import RoborockDockType, \
+    STATE_CODE_TO_STATUS
 from .containers import (
     UserData,
     Status,
@@ -35,7 +35,7 @@
     HomeData,
     MultiMapsList,
     SmartWashParameters,
-    RoborockDeviceInfo,
+    RoborockDeviceInfo, WashTowelMode, DustCollectionMode,
 
 )
 from .roborock_queue import RoborockQueue
@@ -76,17 +76,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()
 
@@ -249,7 +249,7 @@ async def _async_response(self, request_id: int, protocol_id: int = 0) -> tuple[
             del self._waiting_queue[request_id]
 
     def _get_payload(
-            self, method: RoborockCommand, params: list = None, secured=False
+        self, method: RoborockCommand, params: list = None, secured=False
     ):
         timestamp = math.floor(time.time())
         request_id = self._id_counter
@@ -276,7 +276,7 @@ def _get_payload(
         return request_id, timestamp, payload
 
     async def send_command(
-            self, device_id: str, method: RoborockCommand, params: list = None
+        self, device_id: str, method: RoborockCommand, params: list = None
     ):
         raise NotImplementedError
 
@@ -323,17 +323,19 @@ async def get_consumable(self, device_id: str) -> Consumable:
         except RoborockTimeout as e:
             _LOGGER.error(e)
 
-    async def get_washing_mode(self, device_id: str) -> RoborockDockWashingModeType:
+    async def get_washing_mode(self, device_id: str) -> WashTowelMode:
         try:
             washing_mode = await self.send_command(device_id, RoborockCommand.GET_WASH_TOWEL_MODE)
-            return WASH_MODE_MAP.get(washing_mode['wash_mode'])
+            if isinstance(washing_mode, dict):
+                return WashTowelMode(washing_mode)
         except RoborockTimeout as e:
             _LOGGER.error(e)
 
-    async def get_dust_collection_mode(self, device_id: str) -> RoborockDockDustCollectionType:
+    async def get_dust_collection_mode(self, device_id: str) -> DustCollectionMode:
         try:
             dust_collection = await self.send_command(device_id, RoborockCommand.GET_DUST_COLLECTION_MODE)
-            return DUST_COLLECTION_MAP.get(dust_collection['mode'])
+            if isinstance(dust_collection, dict):
+                return DustCollectionMode(dust_collection)
         except RoborockTimeout as e:
             _LOGGER.error(e)
 

roborock/containers.py

@@ -1,7 +1,8 @@
 from enum import Enum
 
 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
+    MOP_INTENSITY_CODES, DOCK_ERROR_TO_TEXT, DOCK_TYPE_MAP, RoborockDockType, RoborockDockDustCollectionType, \
+    DUST_COLLECTION_MAP, WASH_MODE_MAP, RoborockDockWashingModeType
 
 
 class UserDataRRiotReferenceField(str, Enum):
@@ -225,8 +226,12 @@ class SmartWashField(str, Enum):
     WASH_INTERVAL = "wash_interval"
 
 
+class DustCollectionField(str, Enum):
+    MODE = "mode"
+
+
 class WashTowelField(str, Enum):
-    Wash_MODE = "wash_mode"
+    WASH_MODE = "wash_mode"
 
 
 class NetworkInfoField(str, Enum):
@@ -1025,6 +1030,20 @@ def wash_interval(self) -> int:
         return self.get(SmartWashField.WASH_INTERVAL)
 
 
+class DustCollectionMode(RoborockBase):
+
+    @property
+    def mode(self) -> RoborockDockDustCollectionType:
+        return DUST_COLLECTION_MAP.get(self.get(DustCollectionField.MODE))
+
+
+class WashTowelMode(RoborockBase):
+
+    @property
+    def wash_mode(self) -> RoborockDockWashingModeType:
+        return WASH_MODE_MAP.get(self.get(WashTowelField.WASH_MODE))
+
+
 class NetworkInfo(RoborockBase):
     @property
     def ssid(self) -> str:

roborock/typing.py

@@ -2,9 +2,8 @@
 
 from enum import Enum
 
-from .code_mappings import RoborockDockDustCollectionType, RoborockDockWashingModeType
 from .containers import Status, CleanSummary, Consumable, \
-    DNDTimer, CleanRecord, SmartWashParameters
+    DNDTimer, CleanRecord, SmartWashParameters, DustCollectionMode, WashTowelMode
 
 
 class RoborockDevicePropField(str, Enum):
@@ -90,11 +89,12 @@ class RoborockCommand(str, Enum):
 
 
 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
+    def __init__(self, dust_collection_mode: DustCollectionMode,
+                 washing_mode_type: WashTowelMode, mop_wash: SmartWashParameters) -> None:
+        self.dust_collection_mode = dust_collection_mode.mode.value
+        self.washing_mode_type = washing_mode_type.wash_mode.value
+        self.mop_wash_interval = mop_wash.wash_interval
+        self.mop_wash_mode = mop_wash.smart_wash
 
 
 class RoborockDeviceProp:

tests/test_api.py

@@ -52,7 +52,7 @@ async def test_request_code():
 async def test_get_home_data():
     rc = RoborockApiClient("sample@gmail.com")
     with patch("roborock.api.RoborockApiClient._get_base_url") as mock_url, patch(
-            "roborock.api.RoborockApiClient._get_header_client_id") as mock_header_client, patch(
+        "roborock.api.RoborockApiClient._get_header_client_id") as mock_header_client, patch(
         "roborock.api.PreparedRequest.request") as mock_prepared_request:
         mock_prepared_request.side_effect = [{'code': 200, 'msg': 'success', 'data': {"rrHomeId": 1}},
                                              {'code': 200, 'success': True, 'result': HOME_DATA_RAW}]
@@ -70,7 +70,8 @@ async def test_get_dust_collection_mode():
     rmc = RoborockMqttClient(UserData(USER_DATA), device_map)
     with patch("roborock.cloud_api.RoborockMqttClient.send_command") as command:
         command.return_value = {"mode": 1}
-        assert await rmc.get_dust_collection_mode(home_data.devices[0].duid) == RoborockDockDustCollectionType.LIGHT
+        dust = await rmc.get_dust_collection_mode(home_data.devices[0].duid)
+        assert dust.mode == RoborockDockDustCollectionType.LIGHT
 
 
 @pytest.mark.asyncio
@@ -92,4 +93,5 @@ async def test_get_washing_mode():
     rmc = RoborockMqttClient(UserData(USER_DATA), device_map)
     with patch("roborock.cloud_api.RoborockMqttClient.send_command") as command:
         command.return_value = {'wash_mode': 2}
-        assert await rmc.get_washing_mode(home_data.devices[0].duid) == RoborockDockWashingModeType.DEEP
+        washing_mode = await rmc.get_washing_mode(home_data.devices[0].duid)
+        assert washing_mode.wash_mode == RoborockDockWashingModeType.DEEP