feat: extending device status by device model (#51)

* feat: extending device status by device model

* chore: linting

Author: humbertogontijo

roborock/api.py

@@ -13,7 +13,7 @@
 import struct
 import time
 from random import randint
-from typing import Any, Callable, Coroutine, Optional
+from typing import Any, Callable, Coroutine, Optional, Type
 
 import aiohttp
 
@@ -25,6 +25,7 @@
     DNDTimer,
     DustCollectionMode,
     HomeData,
+    ModelStatus,
     MultiMapsList,
     NetworkInfo,
     RoborockDeviceInfo,
@@ -221,9 +222,8 @@ async def send_command(self, method: RoborockCommand, params: Optional[list | di
     async def get_status(self) -> Status | None:
         status = await self.send_command(RoborockCommand.GET_STATUS)
         if isinstance(status, dict):
-            status = Status.from_dict(status)
-            status.update_status(self.device_info.model_specification)
-            return status
+            _cls: Type[Status] = ModelStatus[self.device_info.model]
+            return _cls.from_dict(status)
 
         return None
 

roborock/cli.py

@@ -125,17 +125,13 @@ async def command(ctx, cmd, device_id, params):
     device = next((device for device in devices if device.duid == device_id), None)
     if device is None:
         raise RoborockException("No device found")
-    model_specification = next(
-        (
-            product.model_specification
-            for product in home_data.products
-            if device is not None and product.did == device.duid
-        ),
+    model = next(
+        (product.model for product in home_data.products if device is not None and product.did == device.duid),
         None,
     )
-    if model_specification is None:
-        raise RoborockException(f"Could not find model specifications for device {device.name}")
-    device_info = RoborockDeviceInfo(device=device, model_specification=model_specification)
+    if model is None:
+        raise RoborockException(f"Could not find model for device {device.name}")
+    device_info = RoborockDeviceInfo(device=device, model=model)
     mqtt_client = RoborockMqttClient(login_data.user_data, device_info)
     await mqtt_client.send_command(cmd, params)
     mqtt_client.__del__()

roborock/code_mappings.py

@@ -1,20 +1,9 @@
 from __future__ import annotations
 
 import logging
-from dataclasses import dataclass
 from enum import IntEnum
 from typing import Type
 
-from roborock.const import (
-    ROBOROCK_Q7_MAX,
-    ROBOROCK_S5_MAX,
-    ROBOROCK_S6_MAXV,
-    ROBOROCK_S6_PURE,
-    ROBOROCK_S7,
-    ROBOROCK_S7_MAXV,
-    ROBOROCK_S8_PRO_ULTRA,
-)
-
 _LOGGER = logging.getLogger(__name__)
 
 
@@ -256,65 +245,3 @@ class RoborockDockWashTowelModeCode(RoborockEnum):
     light = 0
     balanced = 1
     deep = 2
-
-
-@dataclass
-class ModelSpecification:
-    model_name: str
-    model_code: str
-    fan_power_code: Type[RoborockFanPowerCode]
-    mop_mode_code: Type[RoborockMopModeCode] | None
-    mop_intensity_code: Type[RoborockMopIntensityCode] | None
-
-
-model_specifications = {
-    ROBOROCK_S5_MAX: ModelSpecification(
-        model_name="Roborock S5 Max",
-        model_code=ROBOROCK_S5_MAX,
-        fan_power_code=RoborockFanSpeedS6Pure,
-        mop_mode_code=None,
-        mop_intensity_code=RoborockMopIntensityV2,
-    ),
-    ROBOROCK_Q7_MAX: ModelSpecification(
-        model_name="Roborock Q7 Max",
-        model_code=ROBOROCK_Q7_MAX,
-        fan_power_code=RoborockFanSpeedQ7Max,
-        mop_mode_code=None,
-        mop_intensity_code=RoborockMopIntensityV2,
-    ),
-    ROBOROCK_S6_MAXV: ModelSpecification(
-        model_name="Roborock S6 MaxV",
-        model_code=ROBOROCK_S6_MAXV,
-        fan_power_code=RoborockFanSpeedE2,
-        mop_mode_code=None,
-        mop_intensity_code=RoborockMopIntensityV2,
-    ),
-    ROBOROCK_S6_PURE: ModelSpecification(
-        model_name="Roborock S6 Pure",
-        model_code=ROBOROCK_S6_PURE,
-        fan_power_code=RoborockFanSpeedS6Pure,
-        mop_mode_code=None,
-        mop_intensity_code=None,
-    ),
-    ROBOROCK_S7_MAXV: ModelSpecification(
-        model_name="Roborock S7 MaxV",
-        model_code=ROBOROCK_S7_MAXV,
-        fan_power_code=RoborockFanSpeedS7MaxV,
-        mop_mode_code=RoborockMopModeS7,
-        mop_intensity_code=RoborockMopIntensityS7,
-    ),
-    ROBOROCK_S7: ModelSpecification(
-        model_name="Roborock S7",
-        model_code=ROBOROCK_S7,
-        fan_power_code=RoborockFanSpeedS7,
-        mop_mode_code=RoborockMopModeS7,
-        mop_intensity_code=RoborockMopIntensityS7,
-    ),
-    ROBOROCK_S8_PRO_ULTRA: ModelSpecification(
-        model_name="Roborock S8 Pro Ultra",
-        model_code=ROBOROCK_S8_PRO_ULTRA,
-        fan_power_code=RoborockFanSpeedS7MaxV,
-        mop_mode_code=RoborockMopModeS8ProUltra,
-        mop_intensity_code=RoborockMopIntensityS7,
-    ),
-}

roborock/containers.py

@@ -9,22 +9,35 @@
 from dacite import Config, from_dict
 
 from .code_mappings import (
-    ModelSpecification,
     RoborockDockDustCollectionModeCode,
     RoborockDockErrorCode,
     RoborockDockTypeCode,
     RoborockDockWashTowelModeCode,
     RoborockErrorCode,
     RoborockFanPowerCode,
+    RoborockFanSpeedE2,
+    RoborockFanSpeedQ7Max,
+    RoborockFanSpeedS6Pure,
+    RoborockFanSpeedS7,
+    RoborockFanSpeedS7MaxV,
     RoborockMopIntensityCode,
+    RoborockMopIntensityS7,
+    RoborockMopIntensityV2,
     RoborockMopModeCode,
+    RoborockMopModeS7,
+    RoborockMopModeS8ProUltra,
     RoborockStateCode,
-    model_specifications,
 )
 from .const import (
     FILTER_REPLACE_TIME,
     MAIN_BRUSH_REPLACE_TIME,
+    ROBOROCK_Q7_MAX,
+    ROBOROCK_S5_MAX,
+    ROBOROCK_S6_MAXV,
+    ROBOROCK_S6_PURE,
+    ROBOROCK_S7,
     ROBOROCK_S7_MAXV,
+    ROBOROCK_S8_PRO_ULTRA,
     SENSOR_DIRTY_REPLACE_TIME,
     SIDE_BRUSH_REPLACE_TIME,
 )
@@ -126,28 +139,6 @@ class HomeDataProduct(RoborockBase):
     capability: Optional[int] = None
     category: Optional[str] = None
     schema: Optional[list[HomeDataProductSchema]] = None
-    model_specification: Optional[ModelSpecification] = None
-
-    def __post_init__(self):
-        if self.model not in model_specifications:
-            _LOGGER.warning("We don't have specific device information for your model, please open an issue.")
-            self.model_specification = model_specifications.get(ROBOROCK_S7_MAXV)
-        else:
-            self.model_specification = model_specifications.get(self.model)
-
-
-@dataclass
-class HomeDataDeviceStatus(RoborockBase):
-    id: Optional[Any] = None
-    name: Optional[Any] = None
-    code: Optional[Any] = None
-    model: Optional[str] = None
-    icon_url: Optional[Any] = None
-    attribute: Optional[Any] = None
-    capability: Optional[Any] = None
-    category: Optional[Any] = None
-    schema: Optional[Any] = None
-    model_specification: Optional[ModelSpecification] = None
 
 
 @dataclass
@@ -175,7 +166,7 @@ class HomeDataDevice(RoborockBase):
     sn: Optional[str] = None
     feature_set: Optional[str] = None
     new_feature_set: Optional[str] = None
-    device_status: Optional[HomeDataDeviceStatus] = None
+    device_status: Optional[dict] = None
     silent_ota_switch: Optional[bool] = None
 
 
@@ -231,14 +222,12 @@ class Status(RoborockBase):
     back_type: Optional[int] = None
     wash_phase: Optional[int] = None
     wash_ready: Optional[int] = None
-    fan_power: Optional[int] = None
-    fan_power_enum: Optional[Type[RoborockFanPowerCode]] = None
+    fan_power: Optional[RoborockFanPowerCode] = None
     dnd_enabled: Optional[int] = None
     map_status: Optional[int] = None
     is_locating: Optional[int] = None
     lock_status: Optional[int] = None
-    water_box_mode: Optional[int] = None
-    water_box_mode_enum: Optional[Type[RoborockMopIntensityCode]] = None
+    water_box_mode: Optional[RoborockMopIntensityCode] = None
     water_box_carriage_status: Optional[int] = None
     mop_forbidden_enable: Optional[int] = None
     camera_status: Optional[int] = None
@@ -251,8 +240,7 @@ class Status(RoborockBase):
     dust_collection_status: Optional[int] = None
     auto_dust_collection: Optional[int] = None
     avoid_count: Optional[int] = None
-    mop_mode: Optional[int] = None
-    mop_mode_enum: Optional[Type[RoborockMopModeCode]] = None
+    mop_mode: Optional[RoborockMopModeCode] = None
     debug_mode: Optional[int] = None
     collision_avoid_status: Optional[int] = None
     switch_map_mode: Optional[int] = None
@@ -261,12 +249,60 @@ class Status(RoborockBase):
     unsave_map_reason: Optional[int] = None
     unsave_map_flag: Optional[int] = None
 
-    def update_status(self, model_specification: ModelSpecification) -> None:
-        self.fan_power_enum = model_specification.fan_power_code.as_enum_dict()[self.fan_power]
-        if model_specification.mop_mode_code is not None:
-            self.mop_mode_enum = model_specification.mop_mode_code.as_enum_dict()[self.mop_mode]
-        if model_specification.mop_intensity_code is not None:
-            self.water_box_mode_enum = model_specification.mop_intensity_code.as_enum_dict()[self.water_box_mode]
+
+@dataclass
+class S5MaxStatus(Status):
+    fan_power: Optional[RoborockFanSpeedS6Pure] = None
+    water_box_mode: Optional[RoborockMopIntensityV2] = None
+
+
+@dataclass
+class Q7MaxStatus(Status):
+    fan_power: Optional[RoborockFanSpeedQ7Max] = None
+    water_box_mode: Optional[RoborockMopIntensityV2] = None
+
+
+@dataclass
+class S6MaxVStatus(Status):
+    fan_power: Optional[RoborockFanSpeedE2] = None
+    water_box_mode: Optional[RoborockMopIntensityV2] = None
+
+
+@dataclass
+class S6PureStatus(Status):
+    fan_power: Optional[RoborockFanSpeedS6Pure] = None
+
+
+@dataclass
+class S7MaxVStatus(Status):
+    fan_power: Optional[RoborockFanSpeedS7MaxV] = None
+    water_box_mode: Optional[RoborockMopIntensityS7] = None
+    mop_mode: Optional[RoborockMopModeS7] = None
+
+
+@dataclass
+class S7Status(Status):
+    fan_power: Optional[RoborockFanSpeedS7] = None
+    water_box_mode: Optional[RoborockMopIntensityS7] = None
+    mop_mode: Optional[RoborockMopModeS7] = None
+
+
+@dataclass
+class S8ProUltraStatus(Status):
+    fan_power: Optional[RoborockFanSpeedS7MaxV] = None
+    water_box_mode: Optional[RoborockMopIntensityS7] = None
+    mop_mode: Optional[RoborockMopModeS8ProUltra] = None
+
+
+ModelStatus: dict[str, Type[Status]] = {
+    ROBOROCK_S5_MAX: S5MaxStatus,
+    ROBOROCK_Q7_MAX: Q7MaxStatus,
+    ROBOROCK_S6_MAXV: S6MaxVStatus,
+    ROBOROCK_S6_PURE: S6PureStatus,
+    ROBOROCK_S7_MAXV: S7MaxVStatus,
+    ROBOROCK_S7: S7Status,
+    ROBOROCK_S8_PRO_ULTRA: S8ProUltraStatus,
+}
 
 
 @dataclass
@@ -385,7 +421,7 @@ class NetworkInfo(RoborockBase):
 @dataclass
 class RoborockDeviceInfo(RoborockBase):
     device: HomeDataDevice
-    model_specification: ModelSpecification
+    model: str
 
 
 @dataclass

tests/conftest.py

@@ -11,7 +11,8 @@ def mqtt_client():
     user_data = UserData.from_dict(USER_DATA)
     home_data = HomeData.from_dict(HOME_DATA_RAW)
     device_info = RoborockDeviceInfo(
-        device=home_data.devices[0], model_specification=home_data.products[0].model_specification
+        device=home_data.devices[0],
+        model=home_data.products[0].model,
     )
     client = RoborockMqttClient(user_data, device_info)
     yield client

tests/test_api.py

@@ -12,7 +12,7 @@
 )
 from roborock.api import PreparedRequest, RoborockApiClient
 from roborock.cloud_api import RoborockMqttClient
-from roborock.containers import RoborockDeviceInfo, Status
+from roborock.containers import RoborockDeviceInfo, S7MaxVStatus
 from tests.mock_data import BASE_URL_REQUEST, GET_CODE_RESPONSE, HOME_DATA_RAW, STATUS, USER_DATA
 
 
@@ -26,9 +26,7 @@ def test_can_create_prepared_request():
 
 def test_can_create_mqtt_roborock():
     home_data = HomeData.from_dict(HOME_DATA_RAW)
-    device_info = RoborockDeviceInfo(
-        device=home_data.devices[0], model_specification=home_data.products[0].model_specification
-    )
+    device_info = RoborockDeviceInfo(device=home_data.devices[0], model=home_data.products[0].model)
     RoborockMqttClient(UserData.from_dict(USER_DATA), device_info)
 
 
@@ -77,9 +75,7 @@ async def test_get_home_data():
 @pytest.mark.asyncio
 async def test_get_dust_collection_mode():
     home_data = HomeData.from_dict(HOME_DATA_RAW)
-    device_info = RoborockDeviceInfo(
-        device=home_data.devices[0], model_specification=home_data.products[0].model_specification
-    )
+    device_info = RoborockDeviceInfo(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 = {"mode": 1}
@@ -91,9 +87,7 @@ async def test_get_dust_collection_mode():
 @pytest.mark.asyncio
 async def test_get_mop_wash_mode():
     home_data = HomeData.from_dict(HOME_DATA_RAW)
-    device_info = RoborockDeviceInfo(
-        device=home_data.devices[0], model_specification=home_data.products[0].model_specification
-    )
+    device_info = RoborockDeviceInfo(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 = {"smart_wash": 0, "wash_interval": 1500}
@@ -106,9 +100,7 @@ async def test_get_mop_wash_mode():
 @pytest.mark.asyncio
 async def test_get_washing_mode():
     home_data = HomeData.from_dict(HOME_DATA_RAW)
-    device_info = RoborockDeviceInfo(
-        device=home_data.devices[0], model_specification=home_data.products[0].model_specification
-    )
+    device_info = RoborockDeviceInfo(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 = {"wash_mode": 2}
@@ -121,15 +113,12 @@ async def test_get_washing_mode():
 @pytest.mark.asyncio
 async def test_get_prop():
     home_data = HomeData.from_dict(HOME_DATA_RAW)
-    device_info = RoborockDeviceInfo(
-        device=home_data.devices[0], model_specification=home_data.products[0].model_specification
-    )
+    device_info = RoborockDeviceInfo(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"):
-        status = Status.from_dict(STATUS)
-        status.update_status(home_data.products[0].model_specification)
+        status = S7MaxVStatus.from_dict(STATUS)
         status.dock_type = RoborockDockTypeCode.auto_empty_dock_pure
         get_status.return_value = status