From 9fe352edc8da51dc6da3d47876b768b4f2e52b4a Mon Sep 17 00:00:00 2001 From: humbertogontijo Date: Thu, 4 May 2023 22:43:55 -0300 Subject: [PATCH 1/2] feat: extending device status by device model --- roborock/api.py | 8 +-- roborock/cli.py | 10 ++-- roborock/code_mappings.py | 73 ------------------------- roborock/containers.py | 112 +++++++++++++++++++++++++------------- tests/conftest.py | 3 +- tests/test_api.py | 25 +++------ tests/test_containers.py | 8 +-- 7 files changed, 95 insertions(+), 144 deletions(-) diff --git a/roborock/api.py b/roborock/api.py index eadef60d..5826db28 100644 --- a/roborock/api.py +++ b/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, + MODEL_STATUS, MultiMapsList, NetworkInfo, RoborockDeviceInfo, @@ -219,9 +220,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] = MODEL_STATUS[self.device_info.model] + return _cls.from_dict(status) return None diff --git a/roborock/cli.py b/roborock/cli.py index 09514413..6962151b 100644 --- a/roborock/cli.py +++ b/roborock/cli.py @@ -125,17 +125,17 @@ 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( + model = next( ( - product.model_specification + 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__() diff --git a/roborock/code_mappings.py b/roborock/code_mappings.py index f8778150..5c31a524 100644 --- a/roborock/code_mappings.py +++ b/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, - ), -} diff --git a/roborock/containers.py b/roborock/containers.py index a951540a..55a83ce1 100644 --- a/roborock/containers.py +++ b/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 + + +MODEL_STATUS: 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 diff --git a/tests/conftest.py b/tests/conftest.py index 7d55bdfe..1c08fbad 100644 --- a/tests/conftest.py +++ b/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 diff --git a/tests/test_api.py b/tests/test_api.py index 7eb0ab5b..3feb8d3d 100644 --- 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 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 diff --git a/tests/test_containers.py b/tests/test_containers.py index a902dc6b..29a21c39 100644 --- a/tests/test_containers.py +++ b/tests/test_containers.py @@ -1,4 +1,4 @@ -from roborock import ROBOROCK_S7_MAXV, CleanRecord, CleanSummary, Consumable, DNDTimer, HomeData, Status, UserData +from roborock import CleanRecord, CleanSummary, Consumable, DNDTimer, HomeData, S7MaxVStatus, UserData from roborock.code_mappings import ( RoborockDockErrorCode, RoborockDockTypeCode, @@ -7,7 +7,6 @@ RoborockMopIntensityS7, RoborockMopModeS7, RoborockStateCode, - model_specifications, ) from .mock_data import CLEAN_RECORD, CLEAN_SUMMARY, CONSUMABLE, DND_TIMER, HOME_DATA_RAW, STATUS, USER_DATA @@ -109,7 +108,7 @@ def test_consumable(): def test_status(): - s = Status.from_dict(STATUS) + s = S7MaxVStatus.from_dict(STATUS) assert s.msg_ver == 2 assert s.msg_seq == 458 assert s.state == RoborockStateCode.charging @@ -152,7 +151,6 @@ def test_status(): assert s.charge_status == 1 assert s.unsave_map_reason == 0 assert s.unsave_map_flag == 0 - s.update_status(model_specification=model_specifications[ROBOROCK_S7_MAXV]) assert s.fan_power == RoborockFanSpeedS7MaxV.balanced assert s.mop_mode == RoborockMopModeS7.standard assert s.water_box_mode == RoborockMopIntensityS7.intense @@ -197,7 +195,7 @@ def test_clean_record(): def test_no_value(): modified_status = STATUS.copy() modified_status["dock_type"] = 9999 - s = Status.from_dict(modified_status) + s = S7MaxVStatus.from_dict(modified_status) assert s.dock_type == RoborockDockTypeCode.missing assert -9999 not in RoborockDockTypeCode.keys() assert "missing" not in RoborockDockTypeCode.values() From 4c71ed221728f036a1526665991dc19b2e8604e6 Mon Sep 17 00:00:00 2001 From: humbertogontijo Date: Thu, 4 May 2023 22:53:39 -0300 Subject: [PATCH 2/2] chore: linting --- roborock/api.py | 4 ++-- roborock/cli.py | 6 +----- roborock/containers.py | 2 +- 3 files changed, 4 insertions(+), 8 deletions(-) diff --git a/roborock/api.py b/roborock/api.py index 5826db28..1c0089d7 100644 --- a/roborock/api.py +++ b/roborock/api.py @@ -25,7 +25,7 @@ DNDTimer, DustCollectionMode, HomeData, - MODEL_STATUS, + ModelStatus, MultiMapsList, NetworkInfo, RoborockDeviceInfo, @@ -220,7 +220,7 @@ 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): - _cls: Type[Status] = MODEL_STATUS[self.device_info.model] + _cls: Type[Status] = ModelStatus[self.device_info.model] return _cls.from_dict(status) return None diff --git a/roborock/cli.py b/roborock/cli.py index 6962151b..ca90c7a7 100644 --- a/roborock/cli.py +++ b/roborock/cli.py @@ -126,11 +126,7 @@ async def command(ctx, cmd, device_id, params): if device is None: raise RoborockException("No device found") model = next( - ( - product.model - for product in home_data.products - if device is not None and product.did == device.duid - ), + (product.model for product in home_data.products if device is not None and product.did == device.duid), None, ) if model is None: diff --git a/roborock/containers.py b/roborock/containers.py index 55a83ce1..078ee5b5 100644 --- a/roborock/containers.py +++ b/roborock/containers.py @@ -294,7 +294,7 @@ class S8ProUltraStatus(Status): mop_mode: Optional[RoborockMopModeS8ProUltra] = None -MODEL_STATUS: dict[str, Type[Status]] = { +ModelStatus: dict[str, Type[Status]] = { ROBOROCK_S5_MAX: S5MaxStatus, ROBOROCK_Q7_MAX: Q7MaxStatus, ROBOROCK_S6_MAXV: S6MaxVStatus,