Fix roborock off peak electricity timer (home-assistant#158292)

Author: allenporter

homeassistant/components/roborock/time.py

@@ -7,7 +7,7 @@
 import logging
 from typing import Any
 
-from roborock.data import DnDTimer
+from roborock.data import DnDTimer, ValleyElectricityTimer
 from roborock.exceptions import RoborockException
 
 from homeassistant.components.time import TimeEntity, TimeEntityDescription
@@ -80,13 +80,14 @@ class RoborockTimeDescription(TimeEntityDescription):
         key="off_peak_start",
         translation_key="off_peak_start",
         trait=lambda api: api.valley_electricity_timer,
-        update_value=lambda trait, desired_time: trait.update_value(
-            [
-                desired_time.hour,
-                desired_time.minute,
-                trait.end_hour,
-                trait.end_minute,
-            ]
+        update_value=lambda trait, desired_time: trait.set_timer(
+            ValleyElectricityTimer(
+                enabled=trait.enabled,
+                start_hour=desired_time.hour,
+                start_minute=desired_time.minute,
+                end_hour=trait.end_hour,
+                end_minute=trait.end_minute,
+            )
         ),
         get_value=lambda trait: datetime.time(
             hour=trait.start_hour, minute=trait.start_minute
@@ -98,13 +99,14 @@ class RoborockTimeDescription(TimeEntityDescription):
         key="off_peak_end",
         translation_key="off_peak_end",
         trait=lambda api: api.valley_electricity_timer,
-        update_value=lambda trait, desired_time: trait.update_value(
-            [
-                trait.start_hour,
-                trait.start_minute,
-                desired_time.hour,
-                desired_time.minute,
-            ]
+        update_value=lambda trait, desired_time: trait.set_timer(
+            ValleyElectricityTimer(
+                enabled=trait.enabled,
+                start_hour=trait.start_hour,
+                start_minute=trait.start_minute,
+                end_hour=desired_time.hour,
+                end_minute=desired_time.minute,
+            )
         ),
         get_value=lambda trait: datetime.time(
             hour=trait.end_hour, minute=trait.end_minute

tests/components/roborock/conftest.py

@@ -21,6 +21,7 @@
     NetworkInfo,
     RoborockBase,
     RoborockDyadStateCode,
+    ValleyElectricityTimer,
     ZeoError,
     ZeoState,
 )
@@ -40,6 +41,9 @@
 from roborock.devices.traits.v1.routines import RoutinesTrait
 from roborock.devices.traits.v1.smart_wash_params import SmartWashParamsTrait
 from roborock.devices.traits.v1.status import StatusTrait
+from roborock.devices.traits.v1.valley_electricity_timer import (
+    ValleyElectricityTimerTrait,
+)
 from roborock.devices.traits.v1.volume import SoundVolumeTrait
 from roborock.devices.traits.v1.wash_towel_mode import WashTowelModeTrait
 from roborock.roborock_message import RoborockDyadDataProtocol, RoborockZeoProtocol
@@ -68,6 +72,7 @@
     STATUS,
     USER_DATA,
     USER_EMAIL,
+    VALLEY_ELECTRICITY_TIMER,
 )
 
 from tests.common import MockConfigEntry
@@ -188,6 +193,25 @@ async def set_dnd_timer(timer: DnDTimer) -> None:
     return dnd_trait
 
 
+def make_valley_electric_timer(dataclass_template: RoborockBase) -> AsyncMock:
+    """Make a function for the fake timer trait that emulates the real behavior."""
+    valley_electric_timer_trait = make_mock_switch(
+        trait_spec=ValleyElectricityTimerTrait,
+        dataclass_template=dataclass_template,
+    )
+
+    async def set_timer(timer: ValleyElectricityTimer) -> None:
+        setattr(valley_electric_timer_trait, "start_hour", timer.start_hour)
+        setattr(valley_electric_timer_trait, "start_minute", timer.start_minute)
+        setattr(valley_electric_timer_trait, "end_hour", timer.end_hour)
+        setattr(valley_electric_timer_trait, "end_minute", timer.end_minute)
+        setattr(valley_electric_timer_trait, "enabled", timer.enabled)
+
+    valley_electric_timer_trait.set_timer = AsyncMock()
+    valley_electric_timer_trait.set_timer.side_effect = set_timer
+    return valley_electric_timer_trait
+
+
 def make_home_trait(
     map_info: list[MultiMapsListMapInfo],
     current_map: int | None,
@@ -257,7 +281,9 @@ def create_v1_properties(network_info: NetworkInfo) -> AsyncMock:
     v1_properties.child_lock = make_mock_switch()
     v1_properties.led_status = make_mock_switch()
     v1_properties.flow_led_status = make_mock_switch()
-    v1_properties.valley_electricity_timer = make_mock_switch()
+    v1_properties.valley_electricity_timer = make_valley_electric_timer(
+        dataclass_template=VALLEY_ELECTRICITY_TIMER,
+    )
     v1_properties.dust_collection_mode = make_mock_trait(
         trait_spec=DustCollectionModeTrait
     )

tests/components/roborock/mock_data.py

@@ -14,6 +14,7 @@
     NetworkInfo,
     S7Status,
     UserData,
+    ValleyElectricityTimer,
 )
 from vacuum_map_parser_base.config.image_config import ImageConfig
 from vacuum_map_parser_base.map_data import ImageData
@@ -1072,6 +1073,16 @@
     }
 )
 
+VALLEY_ELECTRICITY_TIMER = ValleyElectricityTimer.from_dict(
+    {
+        "start_hour": 23,
+        "start_minute": 0,
+        "end_hour": 7,
+        "end_minute": 0,
+        "enabled": 1,
+    }
+)
+
 STATUS = S7Status.from_dict(
     {
         "msg_ver": 2,

tests/components/roborock/test_time.py

@@ -1,10 +1,12 @@
 """Test Roborock Time platform."""
 
+from collections.abc import Callable
 from datetime import time
+from typing import Any
 
 import pytest
 import roborock
-from roborock.data import DnDTimer
+from roborock.data import DnDTimer, RoborockBaseTimer, ValleyElectricityTimer
 
 from homeassistant.components.time import SERVICE_SET_VALUE
 from homeassistant.const import Platform
@@ -22,23 +24,44 @@ def platforms() -> list[Platform]:
     return [Platform.TIME]
 
 
+@pytest.mark.usefixtures("entity_registry_enabled_by_default")
 @pytest.mark.parametrize(
-    ("entity_id", "start_state", "expected_args", "end_state"),
+    ("entity_id", "start_state", "expected_call", "expected_args", "end_state"),
     [
         (
             "time.roborock_s7_maxv_do_not_disturb_begin",
             "22:00:00",
+            lambda x: x.v1_properties.dnd.set_dnd_timer,
             DnDTimer(start_hour=1, start_minute=1, end_hour=7, end_minute=0, enabled=1),
             "01:01:00",
         ),
         (
             "time.roborock_s7_maxv_do_not_disturb_end",
             "07:00:00",
+            lambda x: x.v1_properties.dnd.set_dnd_timer,
             DnDTimer(
                 start_hour=22, start_minute=0, end_hour=1, end_minute=1, enabled=1
             ),
             "01:01:00",
         ),
+        (
+            "time.roborock_s7_maxv_off_peak_start",
+            "23:00:00",
+            lambda x: x.v1_properties.valley_electricity_timer.set_timer,
+            ValleyElectricityTimer(
+                start_hour=1, start_minute=1, end_hour=7, end_minute=0, enabled=1
+            ),
+            "01:01:00",
+        ),
+        (
+            "time.roborock_s7_maxv_off_peak_end",
+            "07:00:00",
+            lambda x: x.v1_properties.valley_electricity_timer.set_timer,
+            ValleyElectricityTimer(
+                start_hour=23, start_minute=0, end_hour=1, end_minute=1, enabled=1
+            ),
+            "01:01:00",
+        ),
     ],
 )
 async def test_update_success(
@@ -48,7 +71,8 @@ async def test_update_success(
     entity_id: str,
     start_state: str,
     end_state: str,
-    expected_args: DnDTimer,
+    expected_call: Callable[[FakeDevice], Any],
+    expected_args: RoborockBaseTimer,
 ) -> None:
     """Test turning switch entities on and off."""
     # Ensure that the entity exist, as these test can pass even if there is no entity.
@@ -64,10 +88,11 @@ async def test_update_success(
         target={"entity_id": entity_id},
     )
 
-    assert fake_vacuum.v1_properties.dnd.set_dnd_timer.call_count == 1
+    call = expected_call(fake_vacuum)
+    assert call.call_count == 1
     # Since we update the begin or end time separately: Verify that the args are built properly
     # by reading the existing value and only updating the relevant fields.
-    assert fake_vacuum.v1_properties.dnd.set_dnd_timer.call_args == ((expected_args,),)
+    assert call.call_args == ((expected_args,),)
 
     state = hass.states.get(entity_id)
     assert state is not None