fix: bug with dock commands

roborock/api.py

@@ -446,11 +446,11 @@ async def get_consumable(self, device_id: str) -> 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)
+        return WASH_MODE_MAP.get(washing_mode['wash_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)
+        return DUST_COLLECTION_MAP.get(dust_collection['mode'])
 
     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)

roborock/code_mappings.py

@@ -14,9 +14,8 @@ class RoborockDockType(str, Enum):
 
 class RoborockDockDustCollectionType(str, Enum):
     SMART = "smart"
-    QUICK = "quick"
-    DAILY = "daily"
-    STRONG = "strong"
+    LIGHT = "light"
+    BALANCED = "balanced"
     MAX = "max"
 
 
@@ -120,9 +119,8 @@ class RoborockDockWashingModeType(str, Enum):
 
 DUST_COLLECTION_MAP = {
     0: RoborockDockDustCollectionType.SMART,
-    1: RoborockDockDustCollectionType.QUICK,
-    2: RoborockDockDustCollectionType.DAILY,
-    3: RoborockDockDustCollectionType.STRONG,
+    1: RoborockDockDustCollectionType.LIGHT,
+    2: RoborockDockDustCollectionType.BALANCED,
     4: RoborockDockDustCollectionType.MAX,
 }
 

tests/test_api.py

@@ -3,7 +3,8 @@
 import paho.mqtt.client as mqtt
 import pytest
 
-from roborock import RoborockClient, UserData, HomeData
+from roborock import RoborockClient, UserData, HomeData, RoborockDockDustCollectionType, SmartWashParameters, \
+    RoborockDockWashingModeType
 from roborock.api import PreparedRequest, RoborockMqttClient
 from tests.mock_data import BASE_URL_REQUEST, GET_CODE_RESPONSE, USER_DATA, HOME_DATA_RAW
 
@@ -50,7 +51,7 @@ async def test_request_code():
 async def test_get_home_data():
     rc = RoborockClient("sample@gmail.com")
     with patch("roborock.api.RoborockClient._get_base_url") as mock_url, patch(
-        "roborock.api.RoborockClient._get_header_client_id") as mock_header_client, patch(
+            "roborock.api.RoborockClient._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}]
@@ -59,3 +60,35 @@ async def test_get_home_data():
         result = await rc.get_home_data(user_data)
 
         assert result == HomeData(HOME_DATA_RAW)
+
+
+@pytest.mark.asyncio
+async def test_get_dust_collection_mode():
+    home_data = HomeData(HOME_DATA_RAW)
+    device_map = {home_data.devices[0].duid: home_data.devices[0]}
+    rmc = RoborockMqttClient(UserData(USER_DATA), device_map)
+    with patch("roborock.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
+
+
+@pytest.mark.asyncio
+async def test_get_mop_wash_mode():
+    home_data = HomeData(HOME_DATA_RAW)
+    device_map = {home_data.devices[0].duid: home_data.devices[0]}
+    rmc = RoborockMqttClient(UserData(USER_DATA), device_map)
+    with patch("roborock.api.RoborockMqttClient.send_command") as command:
+        command.return_value = {'smart_wash': 0, 'wash_interval': 1500}
+        mop_wash = await rmc.get_mop_wash_mode(home_data.devices[0].duid)
+        assert mop_wash.smart_wash == 0
+        assert mop_wash.wash_interval == 1500
+
+
+@pytest.mark.asyncio
+async def test_get_washing_mode():
+    home_data = HomeData(HOME_DATA_RAW)
+    device_map = {home_data.devices[0].duid: home_data.devices[0]}
+    rmc = RoborockMqttClient(UserData(USER_DATA), device_map)
+    with patch("roborock.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