Ensure traits to always reflect the the status of commands

roborock/devices/traits/v1/child_lock.py

@@ -19,7 +19,11 @@ def is_on(self) -> bool:
     async def enable(self) -> None:
         """Enable the child lock."""
         await self.rpc_channel.send_command(RoborockCommand.SET_CHILD_LOCK_STATUS, params={_STATUS_PARAM: 1})
+        # Optimistcally update state to avoid an extra refresh
+        self.lock_status = 1
 
     async def disable(self) -> None:
         """Disable the child lock."""
         await self.rpc_channel.send_command(RoborockCommand.SET_CHILD_LOCK_STATUS, params={_STATUS_PARAM: 0})
+        # Optimistcally update state to avoid an extra refresh
+        self.lock_status = 0

roborock/devices/traits/v1/command.py

@@ -16,7 +16,13 @@ def __post_init__(self) -> None:
         self._rpc_channel = None
 
     async def send(self, command: RoborockCommand | str, params: ParamsType = None) -> Any:
-        """Send a command to the device."""
+        """Send a command to the device.
+
+        Sending a raw command to the device using this method does not update
+        the internal state of any other traits. It is the responsibility of the
+        caller to ensure that any traits affected by the command are refreshed
+        as needed.
+        """
         if not self._rpc_channel:
             raise ValueError("Device trait in invalid state")
         return await self._rpc_channel.send_command(command, params=params)

roborock/devices/traits/v1/common.py

@@ -28,8 +28,14 @@ class V1TraitMixin(ABC):
     Each trait subclass must define a class variable `command` that specifies
     the RoborockCommand used to fetch the trait data from the device. The
     `refresh()` method can be called to update the contents of the trait data
-    from the device. A trait can also support additional commands for updating
-    state associated with the trait.
+    from the device.
+
+    A trait can also support additional commands for updating state associated
+    with the trait. It is expected that a trait will update its own internal
+    state either reflecting the change optimistically or by refreshing the
+    trait state from the device. In cases where one trait caches data that is
+    also represented in another trait, it is the responsibility of the caller
+    to ensure that both traits are refreshed as needed to keep them in sync.
 
     The traits typically subclass RoborockBase to provide serialization
     and deserialization functionality, but this is not strictly required.

roborock/devices/traits/v1/consumeable.py

@@ -45,3 +45,4 @@ class ConsumableTrait(Consumable, common.V1TraitMixin):
     async def reset_consumable(self, consumable: ConsumableAttribute) -> None:
         """Reset a specific consumable attribute on the device."""
         await self.rpc_channel.send_command(RoborockCommand.RESET_CONSUMABLE, params=[consumable.value])
+        await self.refresh()

roborock/devices/traits/v1/do_not_disturb.py

@@ -18,18 +18,24 @@ def is_on(self) -> bool:
     async def set_dnd_timer(self, dnd_timer: DnDTimer) -> None:
         """Set the Do Not Disturb (DND) timer settings of the device."""
         await self.rpc_channel.send_command(RoborockCommand.SET_DND_TIMER, params=dnd_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)
+        await self.refresh()
 
     async def enable(self) -> None:
         """Set the Do Not Disturb (DND) timer settings of the device."""
         await self.rpc_channel.send_command(
             RoborockCommand.SET_DND_TIMER,
             params=self.as_list(),
         )
+        # Optimistcally update state to avoid an extra refresh
+        self.enabled = 1
 
     async def disable(self) -> None:
         """Disable the Do Not Disturb (DND) timer settings of the device."""
         await self.rpc_channel.send_command(RoborockCommand.CLOSE_DND_TIMER)
+        # Optimistcally update state to avoid an extra refresh
+        self.enabled = 0

roborock/devices/traits/v1/flow_led_status.py

@@ -19,7 +19,11 @@ def is_on(self) -> bool:
     async def enable(self) -> None:
         """Enable the Flow LED status."""
         await self.rpc_channel.send_command(RoborockCommand.SET_FLOW_LED_STATUS, params={_STATUS_PARAM: 1})
+        # Optimistcally update state to avoid an extra refresh
+        self.status = 1
 
     async def disable(self) -> None:
         """Disable the Flow LED status."""
         await self.rpc_channel.send_command(RoborockCommand.SET_FLOW_LED_STATUS, params={_STATUS_PARAM: 0})
+        # Optimistcally update state to avoid an extra refresh
+        self.status = 0

roborock/devices/traits/v1/led_status.py

@@ -19,10 +19,14 @@ def is_on(self) -> bool:
     async def enable(self) -> None:
         """Enable the LED status."""
         await self.rpc_channel.send_command(RoborockCommand.SET_LED_STATUS, params=[1])
+        # Optimistcally update state to avoid an extra refresh
+        self.status = 1
 
     async def disable(self) -> None:
         """Disable the LED status."""
         await self.rpc_channel.send_command(RoborockCommand.SET_LED_STATUS, params=[0])
+        # Optimistcally update state to avoid an extra refresh
+        self.status = 0
 
     @classmethod
     def _parse_type_response(cls, response: V1ResponseData) -> LedStatus:

roborock/devices/traits/v1/valley_electricity_timer.py

@@ -19,20 +19,26 @@ def is_on(self) -> bool:
     async def set_timer(self, timer: ValleyElectricityTimer) -> None:
         """Set the Valley Electricity Timer settings of the device."""
         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)
+        await self.refresh()
 
     async def enable(self) -> None:
         """Enable the Valley Electricity Timer settings of the device."""
         await self.rpc_channel.send_command(
             RoborockCommand.SET_VALLEY_ELECTRICITY_TIMER,
             params=self.as_list(),
         )
+        # Optimistcally update state to avoid an extra refresh
+        self.enabled = 1
 
     async def disable(self) -> None:
         """Disable the Valley Electricity Timer settings of the device."""
         await self.rpc_channel.send_command(
             RoborockCommand.CLOSE_VALLEY_ELECTRICITY_TIMER,
         )
+        # Optimistcally update state to avoid an extra refresh
+        self.enabled = 0

roborock/devices/traits/v1/volume.py

@@ -24,3 +24,4 @@ class SoundVolumeTrait(SoundVolume, common.V1TraitMixin):
     async def set_volume(self, volume: int) -> None:
         """Set the sound volume of the device."""
         await self.rpc_channel.send_command(RoborockCommand.CHANGE_SOUND_VOLUME, params=[volume])
+        self.volume = volume

tests/devices/traits/v1/test_consumable.py

@@ -1,6 +1,6 @@
 """Tests for the DoNotDisturbTrait class."""
 
-from unittest.mock import AsyncMock
+from unittest.mock import AsyncMock, call
 
 import pytest
 
@@ -60,11 +60,29 @@ async def test_reset_consumable_data(
     reset_param: str,
 ) -> None:
     """Test successfully resetting consumable data."""
+    mock_rpc_channel.send_command.side_effect = [
+        {},  # Response for RESET_CONSUMABLE
+        # Response for GET_CONSUMABLE after reset
+        {
+            "main_brush_work_time": 5555,
+            "side_brush_work_time": 6666,
+            "filter_work_time": 7777,
+            "filter_element_work_time": 8888,
+            "sensor_dirty_time": 9999,
+        },
+    ]
+
     # Call the method
     await consumable_trait.reset_consumable(consumable)
 
     # Verify the RPC call was made correctly with expected parameters
-    mock_rpc_channel.send_command.assert_called_once_with(RoborockCommand.RESET_CONSUMABLE, params=[reset_param])
-
-
-#
+    assert mock_rpc_channel.send_command.mock_calls == [
+        call(RoborockCommand.RESET_CONSUMABLE, params=[reset_param]),
+        call(RoborockCommand.GET_CONSUMABLE),
+    ]
+    # Verify the consumable data was refreshed correctly
+    assert consumable_trait.main_brush_work_time == 5555
+    assert consumable_trait.side_brush_work_time == 6666
+    assert consumable_trait.filter_work_time == 7777
+    assert consumable_trait.filter_element_work_time == 8888
+    assert consumable_trait.sensor_dirty_time == 9999

tests/devices/traits/v1/test_dnd.py

@@ -1,12 +1,13 @@
 """Tests for the DoNotDisturbTrait class."""
 
-from unittest.mock import AsyncMock
+from unittest.mock import AsyncMock, call
 
 import pytest
 
 from roborock.data import DnDTimer
 from roborock.devices.device import RoborockDevice
 from roborock.devices.traits.v1.do_not_disturb import DoNotDisturbTrait
+from roborock.exceptions import RoborockException
 from roborock.roborock_typing import RoborockCommand
 
 
@@ -74,27 +75,74 @@ async def test_set_dnd_timer_success(
     dnd_trait: DoNotDisturbTrait, mock_rpc_channel: AsyncMock, sample_dnd_timer: DnDTimer
 ) -> None:
     """Test successfully setting DnD timer settings."""
+    mock_rpc_channel.send_command.side_effect = [
+        # Response for SET_DND_TIMER
+        {},
+        # Response for GET_DND_TIMER after updating
+        {
+            "startHour": 22,
+            "startMinute": 0,
+            "endHour": 8,
+            "endMinute": 0,
+            "enabled": 1,
+        },
+    ]
+
     # Call the method
     await dnd_trait.set_dnd_timer(sample_dnd_timer)
 
     # Verify the RPC call was made correctly with dataclass converted to dict
 
     expected_params = [22, 0, 8, 0]
-    mock_rpc_channel.send_command.assert_called_once_with(RoborockCommand.SET_DND_TIMER, params=expected_params)
+    assert mock_rpc_channel.send_command.mock_calls == [
+        call(RoborockCommand.SET_DND_TIMER, params=expected_params),
+        call(RoborockCommand.GET_DND_TIMER),
+    ]
+
+    # Verify the trait state is updated
+    assert dnd_trait.enabled == 1
+    assert dnd_trait.is_on
+    assert dnd_trait.start_hour == 22
+    assert dnd_trait.start_minute == 0
+    assert dnd_trait.end_hour == 8
+    assert dnd_trait.end_minute == 0
 
 
 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 = [
+        # Response for CLOSE_DND_TIMER
+        {},
+        # Response for GET_DND_TIMER after clearing
+        {
+            "startHour": 0,
+            "startMinute": 0,
+            "endHour": 0,
+            "endMinute": 0,
+            "enabled": 0,
+        },
+    ]
+
     # Call the method
     await dnd_trait.clear_dnd_timer()
 
     # Verify the RPC call was made correctly
-    mock_rpc_channel.send_command.assert_called_once_with(RoborockCommand.CLOSE_DND_TIMER)
+    assert mock_rpc_channel.send_command.mock_calls == [
+        call(RoborockCommand.CLOSE_DND_TIMER),
+        call(RoborockCommand.GET_DND_TIMER),
+    ]
+
+    # Verify the trait state is updated
+    assert dnd_trait.enabled == 0
+    assert not dnd_trait.is_on
+    assert dnd_trait.start_hour == 0
+    assert dnd_trait.start_minute == 0
+    assert dnd_trait.end_hour == 0
+    assert dnd_trait.end_minute == 0
 
 
 async def test_get_dnd_timer_propagates_exception(dnd_trait: DoNotDisturbTrait, mock_rpc_channel: AsyncMock) -> None:
     """Test that exceptions from RPC channel are propagated in get_dnd_timer."""
-    from roborock.exceptions import RoborockException
 
     # Setup mock to raise an exception
     mock_rpc_channel.send_command.side_effect = RoborockException("Communication error")