diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 755f4f90..81883851 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -42,7 +42,7 @@ repos: hooks: - id: mypy exclude: cli.py - additional_dependencies: [ "types-paho-mqtt" ] + additional_dependencies: [ "types-paho-mqtt", "types-pyyaml" ] - repo: https://github.com/alessandrojcm/commitlint-pre-commit-hook rev: v9.23.0 hooks: diff --git a/roborock/data/v1/v1_containers.py b/roborock/data/v1/v1_containers.py index cf988590..09050186 100644 --- a/roborock/data/v1/v1_containers.py +++ b/roborock/data/v1/v1_containers.py @@ -103,8 +103,10 @@ class StatusField(FieldNameBase): to understand if a feature is supported by the device using `is_field_supported`. The enum values are names of fields in the `Status` class. Each field is - annotated with `requires_schema_code` metadata to map the field to a schema - code in the product schema, which may have a different name than the field/attribute name. + annotated with one of the following: + - `requires_schema_code` metadata to map the field to a schema code in the + product schema, which may have a different name than the field/attribute name. + - `requires_supported_feature` metadata to map the field to a field in `DeviceFeatures`. """ STATE = "state" @@ -113,20 +115,25 @@ class StatusField(FieldNameBase): WATER_BOX_MODE = "water_box_mode" CHARGE_STATUS = "charge_status" DRY_STATUS = "dry_status" + CLEAN_PERCENT = "clean_percent" -def _requires_schema_code(requires_schema_code: str, default=None) -> Any: +def _requires_schema_code(requires_schema_code: str, default: Any = None) -> Any: return field(metadata={"requires_schema_code": requires_schema_code}, default=default) +def _requires_supported_feature(requires_supported_feature: str, default: Any = None) -> Any: + return field(metadata={"requires_supported_feature": requires_supported_feature}, default=default) + + @dataclass class Status(RoborockBase): """This status will be deprecated in favor of StatusV2.""" msg_ver: int | None = None msg_seq: int | None = None - state: RoborockStateCode | None = _requires_schema_code("state", default=None) - battery: int | None = _requires_schema_code("battery", default=None) + state: RoborockStateCode | None = _requires_schema_code("state") + battery: int | None = _requires_schema_code("battery") clean_time: int | None = None clean_area: int | None = None error_code: RoborockErrorCode | None = None @@ -139,12 +146,12 @@ class Status(RoborockBase): back_type: int | None = None wash_phase: int | None = None wash_ready: int | None = None - fan_power: RoborockFanPowerCode | None = _requires_schema_code("fan_power", default=None) + fan_power: RoborockFanPowerCode | None = _requires_schema_code("fan_power") dnd_enabled: int | None = None map_status: int | None = None is_locating: int | None = None lock_status: int | None = None - water_box_mode: RoborockMopIntensityCode | None = _requires_schema_code("water_box_mode", default=None) + water_box_mode: RoborockMopIntensityCode | None = _requires_schema_code("water_box_mode") water_box_carriage_status: int | None = None mop_forbidden_enable: int | None = None camera_status: int | None = None @@ -162,15 +169,15 @@ class Status(RoborockBase): collision_avoid_status: int | None = None switch_map_mode: int | None = None dock_error_status: RoborockDockErrorCode | None = None - charge_status: int | None = _requires_schema_code("charge_status", default=None) + charge_status: int | None = _requires_schema_code("charge_status") unsave_map_reason: int | None = None unsave_map_flag: int | None = None wash_status: int | None = None distance_off: int | None = None in_warmup: int | None = None - dry_status: int | None = _requires_schema_code("drying_status", default=None) + dry_status: int | None = _requires_schema_code("drying_status") rdt: int | None = None - clean_percent: int | None = None + clean_percent: int | None = _requires_supported_feature("is_support_clean_estimate") rss: int | None = None dss: int | None = None common_status: int | None = None diff --git a/roborock/device_features.py b/roborock/device_features.py index b8ecf483..ef3ea19f 100644 --- a/roborock/device_features.py +++ b/roborock/device_features.py @@ -1,8 +1,6 @@ -from __future__ import annotations - from dataclasses import dataclass, field, fields from enum import IntEnum, StrEnum -from typing import Any +from typing import Any, Self from roborock.data.code_mappings import RoborockProductNickname from roborock.data.containers import RoborockBase @@ -566,7 +564,7 @@ def from_feature_flags( new_feature_info_str: str, feature_info: list[int], product_nickname: RoborockProductNickname | None, - ) -> DeviceFeatures: + ) -> Self: """Creates a DeviceFeatures instance from raw feature flags. :param new_feature_info: A int from get_init_status (sometimes can be found in homedata, but it is not always) :param new_feature_info_str: A hex string from get_init_status or home_data. diff --git a/roborock/devices/traits/v1/device_features.py b/roborock/devices/traits/v1/device_features.py index ed6e186e..87971643 100644 --- a/roborock/devices/traits/v1/device_features.py +++ b/roborock/devices/traits/v1/device_features.py @@ -40,10 +40,11 @@ def __init__(self, product: HomeDataProduct, device_cache: DeviceCache) -> None: self.converter = DeviceTraitsConverter(product) self._product = product self._device_cache = device_cache - # All fields of DeviceFeatures are required. Initialize them to False + # All boolean fields of DeviceFeatures are required. Initialize them to False # so we have some known state. for field in fields(self): - setattr(self, field.name, False) + if field.type is bool: + setattr(self, field.name, False) def is_field_supported(self, cls: type[RoborockBase], field_name: FieldNameBase) -> bool: """Determines if the specified field is supported by this device. @@ -61,11 +62,12 @@ def is_field_supported(self, cls: type[RoborockBase], field_name: FieldNameBase) raise ValueError(f"Field {field_name} not found in {cls}") requires_schema_code = dataclass_field.metadata.get("requires_schema_code", None) - if requires_schema_code is None: - # We assume the field is supported - return True - # If the field requires a protocol that is not supported, we return False - return requires_schema_code in self._product.supported_schema_codes + if requires_schema_code is not None: + return requires_schema_code in self._product.supported_schema_codes + requires_supported_feature = dataclass_field.metadata.get("requires_supported_feature", None) + if requires_supported_feature is not None: + return getattr(self, requires_supported_feature) + return True async def refresh(self) -> None: """Refresh the contents of this trait. diff --git a/tests/devices/traits/v1/__snapshots__/test_device_features.ambr b/tests/devices/traits/v1/__snapshots__/test_device_features.ambr index 9b2a4827..a15f714c 100644 --- a/tests/devices/traits/v1/__snapshots__/test_device_features.ambr +++ b/tests/devices/traits/v1/__snapshots__/test_device_features.ambr @@ -3,6 +3,7 @@ dict({ 'battery': True, 'charge_status': True, + 'clean_percent': True, 'dry_status': True, 'fan_power': True, 'state': True, @@ -13,6 +14,7 @@ dict({ 'battery': True, 'charge_status': True, + 'clean_percent': True, 'dry_status': True, 'fan_power': True, 'state': True, @@ -23,6 +25,7 @@ dict({ 'battery': True, 'charge_status': True, + 'clean_percent': True, 'dry_status': True, 'fan_power': True, 'state': True, diff --git a/tests/devices/traits/v1/fixtures.py b/tests/devices/traits/v1/fixtures.py index 08397493..384ebdab 100644 --- a/tests/devices/traits/v1/fixtures.py +++ b/tests/devices/traits/v1/fixtures.py @@ -1,6 +1,7 @@ """Fixtures for V1 trait tests.""" from copy import deepcopy +from typing import Any from unittest.mock import AsyncMock import pytest @@ -107,16 +108,34 @@ def dock_type_code_fixture(request: pytest.FixtureRequest) -> RoborockDockTypeCo return RoborockDockTypeCode.s7_max_ultra_dock +@pytest.fixture(name="mock_app_get_init_status") +def mock_app_get_init_status_fixture(device_info: HomeDataDevice, products: list[HomeDataProduct]) -> dict[str, Any]: + """Fixture to provide model-specific APP_GET_INIT_STATUS data. + + Uses real device feature data from device_info.yaml when available for the + product model, falling back to the default mock data otherwise. + """ + product = next((product for product in products if product.id == device_info.product_id), None) + if product is None: + raise ValueError(f"Product {device_info.product_id} not found") + device_info_data = mock_data.DEVICE_INFO.get(product.model, {}) + return { + **mock_data.APP_GET_INIT_STATUS, + **device_info_data, + } + + @pytest.fixture(autouse=True) async def discover_features_fixture( device: RoborockDevice, mock_rpc_channel: AsyncMock, + mock_app_get_init_status: dict[str, Any], dock_type_code: RoborockDockTypeCode | None, ) -> None: """Fixture to handle device feature discovery.""" assert device.v1_properties mock_rpc_channel.send_command.side_effect = [ - [mock_data.APP_GET_INIT_STATUS], + [mock_app_get_init_status], { **mock_data.STATUS, "dock_type": dock_type_code, diff --git a/tests/mock_data.py b/tests/mock_data.py index 73292c3b..abfbd650 100644 --- a/tests/mock_data.py +++ b/tests/mock_data.py @@ -5,6 +5,8 @@ import pathlib from typing import Any +import yaml + # All data is based on a U.S. customer with a Roborock S7 MaxV Ultra USER_EMAIL = "user@domain.com" @@ -141,6 +143,18 @@ ZEO_ONE_DEVICE_DATA = DEVICES["home_data_device_zeo_one.json"] SAROS_10R_DEVICE_DATA = DEVICES["home_data_device_saros_10r.json"] +# Additional Device Features info from YAML keyed by product model. +# Each entry contains the fields needed for APP_GET_INIT_STATUS responses. +_DEVICE_INFO_DATA = yaml.safe_load(pathlib.Path("device_info.yaml").read_text()) +DEVICE_INFO: dict[str, dict[str, Any]] = { + product_model: { + "new_feature_info": data.get("new_feature_info"), + "new_feature_info_str": data.get("new_feature_info_str"), + "feature_info": data.get("feature_info", []), + } + for product_model, data in _DEVICE_INFO_DATA.items() +} + HOME_DATA_RAW: dict[str, Any] = { "id": 123456,