feat: add set_start_time and set_end_time to V1 timer traits

roborock/devices/traits/v1/do_not_disturb.py

@@ -1,3 +1,5 @@
+import datetime
+
 from roborock.data import DnDTimer
 from roborock.devices.traits.v1 import common
 from roborock.roborock_typing import RoborockCommand
@@ -21,6 +23,30 @@ async def set_dnd_timer(self, dnd_timer: DnDTimer) -> None:
         await self.rpc_channel.send_command(RoborockCommand.SET_DND_TIMER, params=dnd_timer.as_list())
         await self.refresh()
 
+    async def set_start_time(self, start_time: datetime.time) -> None:
+        """Set the start time of the Do Not Disturb (DND) timer."""
+        timer = DnDTimer(
+            start_hour=start_time.hour,
+            start_minute=start_time.minute,
+            end_hour=self.end_hour,
+            end_minute=self.end_minute,
+            enabled=self.enabled,
+        )
+        await self.rpc_channel.send_command(RoborockCommand.SET_DND_TIMER, params=timer.as_list())
+        await self.refresh()
+
+    async def set_end_time(self, end_time: datetime.time) -> None:
+        """Set the end time of the Do Not Disturb (DND) timer."""
+        timer = DnDTimer(
+            start_hour=self.start_hour,
+            start_minute=self.start_minute,
+            end_hour=end_time.hour,
+            end_minute=end_time.minute,
+            enabled=self.enabled,
+        )
+        await self.rpc_channel.send_command(RoborockCommand.SET_DND_TIMER, params=timer.as_list())
+        await self.refresh()
+
     async def clear_dnd_timer(self) -> None:
         """Clear the Do Not Disturb (DND) timer settings of the device."""
         await self.rpc_channel.send_command(RoborockCommand.CLOSE_DND_TIMER)

roborock/devices/traits/v1/valley_electricity_timer.py

@@ -1,3 +1,5 @@
+import datetime
+
 from roborock.data import ValleyElectricityTimer
 from roborock.devices.traits.v1 import common
 from roborock.roborock_typing import RoborockCommand
@@ -22,6 +24,30 @@ async def set_timer(self, timer: ValleyElectricityTimer) -> None:
         await self.rpc_channel.send_command(RoborockCommand.SET_VALLEY_ELECTRICITY_TIMER, params=timer.as_list())
         await self.refresh()
 
+    async def set_start_time(self, start_time: datetime.time) -> None:
+        """Set the start time of the Valley Electricity Timer."""
+        timer = ValleyElectricityTimer(
+            start_hour=start_time.hour,
+            start_minute=start_time.minute,
+            end_hour=self.end_hour,
+            end_minute=self.end_minute,
+            enabled=self.enabled,
+        )
+        await self.rpc_channel.send_command(RoborockCommand.SET_VALLEY_ELECTRICITY_TIMER, params=timer.as_list())
+        await self.refresh()
+
+    async def set_end_time(self, end_time: datetime.time) -> None:
+        """Set the end time of the Valley Electricity Timer."""
+        timer = ValleyElectricityTimer(
+            start_hour=self.start_hour,
+            start_minute=self.start_minute,
+            end_hour=end_time.hour,
+            end_minute=end_time.minute,
+            enabled=self.enabled,
+        )
+        await self.rpc_channel.send_command(RoborockCommand.SET_VALLEY_ELECTRICITY_TIMER, params=timer.as_list())
+        await self.refresh()
+
     async def clear_timer(self) -> None:
         """Clear the Valley Electricity Timer settings of the device."""
         await self.rpc_channel.send_command(RoborockCommand.CLOSE_VALLEY_ELECTRICITY_TIMER)

tests/devices/traits/v1/test_dnd.py

@@ -1,5 +1,6 @@
 """Tests for the DoNotDisturbTrait class."""
 
+import datetime
 from unittest.mock import AsyncMock, call
 
 import pytest
@@ -48,6 +49,8 @@ async def test_get_dnd_timer_success(
     assert dnd_trait.end_minute == 0
     assert dnd_trait.enabled == 1
     assert dnd_trait.is_on
+    assert dnd_trait.start_time == datetime.time(22, 0)
+    assert dnd_trait.end_time == datetime.time(8, 0)
 
     # Verify the RPC call was made correctly
     mock_rpc_channel.send_command.assert_called_once_with(RoborockCommand.GET_DND_TIMER)
@@ -108,6 +111,113 @@ async def test_set_dnd_timer_success(
     assert dnd_trait.end_minute == 0
 
 
+async def test_set_start_time_success(dnd_trait: DoNotDisturbTrait, mock_rpc_channel: AsyncMock) -> None:
+    """Test successfully setting start time on DND timer."""
+    mock_rpc_channel.send_command.side_effect = [
+        # Response for SET_DND_TIMER
+        {},
+        # Response for GET_DND_TIMER after updating
+        {
+            "startHour": 23,
+            "startMinute": 30,
+            "endHour": 8,
+            "endMinute": 0,
+            "enabled": 1,
+        },
+    ]
+
+    # Set up initial trait state so it has an end time
+    dnd_trait.start_hour = 22
+    dnd_trait.start_minute = 0
+    dnd_trait.end_hour = 8
+    dnd_trait.end_minute = 0
+
+    await dnd_trait.set_start_time(datetime.time(23, 30))
+
+    # Verify the RPC call used the updated start hour/minute but kept the old end hour/minute
+    expected_params = [23, 30, 8, 0]
+    assert mock_rpc_channel.send_command.mock_calls == [
+        call(RoborockCommand.SET_DND_TIMER, params=expected_params),
+        call(RoborockCommand.GET_DND_TIMER),
+    ]
+
+    assert dnd_trait.start_hour == 23
+    assert dnd_trait.start_minute == 30
+    assert dnd_trait.start_time == datetime.time(23, 30)
+
+
+async def test_set_end_time_success(dnd_trait: DoNotDisturbTrait, mock_rpc_channel: AsyncMock) -> None:
+    """Test successfully setting end time on DND timer."""
+    mock_rpc_channel.send_command.side_effect = [
+        # Response for SET_DND_TIMER
+        {},
+        # Response for GET_DND_TIMER after updating
+        {
+            "startHour": 22,
+            "startMinute": 0,
+            "endHour": 9,
+            "endMinute": 45,
+            "enabled": 1,
+        },
+    ]
+
+    # Set up initial trait state so it has a start time
+    dnd_trait.start_hour = 22
+    dnd_trait.start_minute = 0
+    dnd_trait.end_hour = 8
+    dnd_trait.end_minute = 0
+
+    await dnd_trait.set_end_time(datetime.time(9, 45))
+
+    # Verify the RPC call used the old start hour/minute but updated the end hour/minute
+    expected_params = [22, 0, 9, 45]
+    assert mock_rpc_channel.send_command.mock_calls == [
+        call(RoborockCommand.SET_DND_TIMER, params=expected_params),
+        call(RoborockCommand.GET_DND_TIMER),
+    ]
+
+    assert dnd_trait.end_hour == 9
+    assert dnd_trait.end_minute == 45
+    assert dnd_trait.end_time == datetime.time(9, 45)
+
+
+async def test_set_end_time_when_start_time_not_set(dnd_trait: DoNotDisturbTrait, mock_rpc_channel: AsyncMock) -> None:
+    """Test setting end time when start time is not set."""
+    mock_rpc_channel.send_command.side_effect = [
+        # Response for SET_DND_TIMER
+        {},
+        # Response for GET_DND_TIMER after updating
+        {
+            "startHour": None,
+            "startMinute": None,
+            "endHour": 9,
+            "endMinute": 45,
+            "enabled": 1,
+        },
+    ]
+
+    # Set up initial trait state so it does NOT have a start time
+    dnd_trait.start_hour = None
+    dnd_trait.start_minute = None
+    dnd_trait.end_hour = None
+    dnd_trait.end_minute = None
+
+    await dnd_trait.set_end_time(datetime.time(9, 45))
+
+    expected_params = [None, None, 9, 45]
+    assert mock_rpc_channel.send_command.mock_calls == [
+        call(RoborockCommand.SET_DND_TIMER, params=expected_params),
+        call(RoborockCommand.GET_DND_TIMER),
+    ]
+
+    assert dnd_trait.start_hour is None
+    assert dnd_trait.start_minute is None
+    assert dnd_trait.end_hour == 9
+    assert dnd_trait.end_minute == 45
+    assert dnd_trait.end_time == datetime.time(9, 45)
+    assert dnd_trait.start_time is None
+
+
 async def test_clear_dnd_timer_success(dnd_trait: DoNotDisturbTrait, mock_rpc_channel: AsyncMock) -> None:
     """Test successfully clearing DnD timer settings."""
     mock_rpc_channel.send_command.side_effect = [