feat: update schema requirements in StatusV2 and Consumable (#806)

* feat: update schema requirements in StatusV2

Update the StatusV2 to check the schema for supported fields.

* fix: update error_code and other fields to use _requires_schema_code for schema validation

* feat: Add consumable fields

* refactor: remove redundant default=None arguments from _requires_schema_code calls in v1_containers

Author: allenporter

roborock/data/v1/v1_containers.py

@@ -113,6 +113,7 @@ class StatusField(FieldNameBase):
     WATER_BOX_MODE = "water_box_mode"
     CHARGE_STATUS = "charge_status"
     DRY_STATUS = "dry_status"
+    ERROR_CODE = "error_code"
 
 
 def _requires_schema_code(requires_schema_code: str, default=None) -> Any:
@@ -125,11 +126,11 @@ class Status(RoborockBase):
 
     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
+    error_code: RoborockErrorCode | None = _requires_schema_code("error_code")
     map_present: int | None = None
     in_cleaning: RoborockInCleaning | None = None
     in_returning: int | None = None
@@ -139,12 +140,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,13 +163,13 @@ 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
     rss: int | None = None
@@ -293,11 +294,11 @@ class StatusV2(RoborockBase):
 
     msg_ver: int | None = None
     msg_seq: int | None = None
-    state: RoborockStateCode | None = None
-    battery: int | None = 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
+    error_code: RoborockErrorCode | None = _requires_schema_code("error_code")
     map_present: int | None = None
     in_cleaning: RoborockInCleaning | None = None
     in_returning: int | None = None
@@ -307,12 +308,12 @@ class StatusV2(RoborockBase):
     back_type: int | None = None
     wash_phase: int | None = None
     wash_ready: int | None = None
-    fan_power: int | None = None
+    fan_power: int | 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: int | None = None
+    water_box_mode: int | 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
@@ -329,14 +330,14 @@ class StatusV2(RoborockBase):
     debug_mode: int | None = None
     collision_avoid_status: int | None = None
     switch_map_mode: int | None = None
-    dock_error_status: RoborockDockErrorCode | None = None
-    charge_status: int | None = None
+    dock_error_status: RoborockDockErrorCode | None = _requires_schema_code("dock_error_status")
+    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 = None
+    dry_status: int | None = _requires_schema_code("drying_status")
     rdt: int | None = None
     clean_percent: int | None = None
     rss: int | None = None
@@ -624,11 +625,27 @@ class CleanSummaryWithDetail(CleanSummary):
     last_clean_record: CleanRecord | None = None
 
 
+class ConsumableField(FieldNameBase):
+    """An enum that represents a field in the `Consumable` class.
+
+    This is used with `roborock.devices.traits.v1.status.DeviceFeaturesTrait`
+    to understand if a feature is supported by the device using `is_field_supported`.
+
+    The enum values are names of fields in the `Consumable` 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.
+    """
+
+    MAIN_BRUSH_WORK_TIME = "main_brush_work_time"
+    SIDE_BRUSH_WORK_TIME = "side_brush_work_time"
+    FILTER_WORK_TIME = "filter_work_time"
+
+
 @dataclass
 class Consumable(RoborockBase):
-    main_brush_work_time: int | None = None
-    side_brush_work_time: int | None = None
-    filter_work_time: int | None = None
+    main_brush_work_time: int | None = field(metadata={"requires_schema_code": "main_brush_life"}, default=None)
+    side_brush_work_time: int | None = field(metadata={"requires_schema_code": "side_brush_life"}, default=None)
+    filter_work_time: int | None = field(metadata={"requires_schema_code": "filter_life"}, default=None)
     filter_element_work_time: int | None = None
     sensor_dirty_time: int | None = None
     strainer_work_times: int | None = None

tests/devices/traits/v1/__snapshots__/test_device_features.ambr

@@ -4,6 +4,7 @@
     'battery': True,
     'charge_status': True,
     'dry_status': True,
+    'error_code': True,
     'fan_power': True,
     'state': True,
     'water_box_mode': True,
@@ -14,6 +15,7 @@
     'battery': True,
     'charge_status': True,
     'dry_status': True,
+    'error_code': True,
     'fan_power': True,
     'state': True,
     'water_box_mode': True,
@@ -24,8 +26,30 @@
     'battery': True,
     'charge_status': True,
     'dry_status': True,
+    'error_code': True,
     'fan_power': True,
     'state': True,
     'water_box_mode': True,
   })
 # ---
+# name: test_is_consumable_field_supported[home_data_device_s5e.json]
+  dict({
+    'filter_work_time': True,
+    'main_brush_work_time': True,
+    'side_brush_work_time': True,
+  })
+# ---
+# name: test_is_consumable_field_supported[home_data_device_s7_maxv.json]
+  dict({
+    'filter_work_time': True,
+    'main_brush_work_time': True,
+    'side_brush_work_time': True,
+  })
+# ---
+# name: test_is_consumable_field_supported[home_data_device_saros_10r.json]
+  dict({
+    'filter_work_time': True,
+    'main_brush_work_time': True,
+    'side_brush_work_time': True,
+  })
+# ---

tests/devices/traits/v1/test_device_features.py

@@ -4,8 +4,9 @@
 from syrupy import SnapshotAssertion
 
 from roborock.data import HomeDataDevice
-from roborock.data.v1.v1_containers import StatusField
+from roborock.data.v1.v1_containers import ConsumableField, StatusField
 from roborock.devices.device import RoborockDevice
+from roborock.devices.traits.v1.consumeable import ConsumableTrait
 from roborock.devices.traits.v1.status import StatusTrait
 from tests import mock_data
 
@@ -29,5 +30,28 @@ async def test_is_attribute_supported(
     assert device.v1_properties.device_features is not None
     device_features_trait = device.v1_properties.device_features
 
-    is_supported = {field.value: device_features_trait.is_field_supported(StatusTrait, field) for field in StatusField}
-    assert is_supported == snapshot
+    is_v1_supported = {
+        field.value: device_features_trait.is_field_supported(StatusTrait, field) for field in StatusField
+    }
+    assert is_v1_supported == snapshot
+
+
+@pytest.mark.parametrize(
+    ("device_info"),
+    V1_DEVICES.values(),
+    ids=list(V1_DEVICES.keys()),
+)
+async def test_is_consumable_field_supported(
+    device_info: HomeDataDevice,
+    device: RoborockDevice,
+    snapshot: SnapshotAssertion,
+) -> None:
+    """Test if a field is supported."""
+    assert device.v1_properties is not None
+    assert device.v1_properties.device_features is not None
+    device_features_trait = device.v1_properties.device_features
+
+    is_v1_supported = {
+        field.value: device_features_trait.is_field_supported(ConsumableTrait, field) for field in ConsumableField
+    }
+    assert is_v1_supported == snapshot