fix: Add a common base type for all switches (#548)

Author: allenporter

roborock/devices/traits/v1/__init__.py

@@ -24,6 +24,7 @@
 from .maps import MapsTrait
 from .rooms import RoomsTrait
 from .status import StatusTrait
+from .valley_electricity_timer import ValleyElectricityTimerTrait
 from .volume import SoundVolumeTrait
 
 _LOGGER = logging.getLogger(__name__)
@@ -44,6 +45,7 @@
     "ChildLockTrait",
     "FlowLedStatusTrait",
     "LedStatusTrait",
+    "ValleyElectricityTimerTrait",
 ]
 
 
@@ -71,6 +73,7 @@ class PropertiesApi(Trait):
     child_lock: ChildLockTrait | None = None
     led_status: LedStatusTrait | None = None
     flow_led_status: FlowLedStatusTrait | None = None
+    valley_electricity_timer: ValleyElectricityTimerTrait | None = None
 
     def __init__(
         self,

roborock/devices/traits/v1/child_lock.py

@@ -5,12 +5,17 @@
 _STATUS_PARAM = "lock_status"
 
 
-class ChildLockTrait(ChildLockStatus, common.V1TraitMixin):
+class ChildLockTrait(ChildLockStatus, common.V1TraitMixin, common.RoborockSwitchBase):
     """Trait for controlling the child lock of a Roborock device."""
 
     command = RoborockCommand.GET_CHILD_LOCK_STATUS
     requires_feature = "is_set_child_supported"
 
+    @property
+    def is_on(self) -> bool:
+        """Return whether the child lock is enabled."""
+        return self.lock_status == 1
+
     async def enable(self) -> None:
         """Enable the child lock."""
         await self.rpc_channel.send_command(RoborockCommand.SET_CHILD_LOCK_STATUS, params={_STATUS_PARAM: 1})

roborock/devices/traits/v1/common.py

@@ -4,7 +4,7 @@
 """
 
 import logging
-from abc import ABC
+from abc import ABC, abstractmethod
 from dataclasses import dataclass, fields
 from typing import ClassVar, Self
 
@@ -124,6 +124,23 @@ def _parse_response(cls, response: V1ResponseData) -> Self:
         return cls(**{value_field: response})
 
 
+class RoborockSwitchBase(ABC):
+    """Base class for traits that represent a boolean switch."""
+
+    @property
+    @abstractmethod
+    def is_on(self) -> bool:
+        """Return whether the switch is on."""
+
+    @abstractmethod
+    async def enable(self) -> None:
+        """Enable the switch."""
+
+    @abstractmethod
+    async def disable(self) -> None:
+        """Disable the switch."""
+
+
 def mqtt_rpc_channel(cls):
     """Decorator to mark a function as cloud only.
 

roborock/devices/traits/v1/do_not_disturb.py

@@ -5,11 +5,16 @@
 _ENABLED_PARAM = "enabled"
 
 
-class DoNotDisturbTrait(DnDTimer, common.V1TraitMixin):
+class DoNotDisturbTrait(DnDTimer, common.V1TraitMixin, common.RoborockSwitchBase):
     """Trait for managing Do Not Disturb (DND) settings on Roborock devices."""
 
     command = RoborockCommand.GET_DND_TIMER
 
+    @property
+    def is_on(self) -> bool:
+        """Return whether the Do Not Disturb (DND) timer is enabled."""
+        return self.enabled == 1
+
     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_dict())

roborock/devices/traits/v1/flow_led_status.py

@@ -5,12 +5,17 @@
 _STATUS_PARAM = "status"
 
 
-class FlowLedStatusTrait(FlowLedStatus, common.V1TraitMixin):
+class FlowLedStatusTrait(FlowLedStatus, common.V1TraitMixin, common.RoborockSwitchBase):
     """Trait for controlling the Flow LED status of a Roborock device."""
 
     command = RoborockCommand.GET_FLOW_LED_STATUS
     requires_feature = "is_flow_led_setting_supported"
 
+    @property
+    def is_on(self) -> bool:
+        """Return whether the Flow LED status is enabled."""
+        return self.status == 1
+
     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})

roborock/devices/traits/v1/led_status.py

@@ -5,12 +5,17 @@
 from .common import V1ResponseData
 
 
-class LedStatusTrait(LedStatus, common.V1TraitMixin):
+class LedStatusTrait(LedStatus, common.V1TraitMixin, common.RoborockSwitchBase):
     """Trait for controlling the LED status of a Roborock device."""
 
     command = RoborockCommand.GET_LED_STATUS
     requires_feature = "is_led_status_switch_supported"
 
+    @property
+    def is_on(self) -> bool:
+        """Return whether the LED status is enabled."""
+        return self.status == 1
+
     async def enable(self) -> None:
         """Enable the LED status."""
         await self.rpc_channel.send_command(RoborockCommand.SET_LED_STATUS, params=[1])

roborock/devices/traits/v1/valley_electricity_timer.py

@@ -11,6 +11,11 @@ class ValleyElectricityTimerTrait(ValleyElectricityTimer, common.V1TraitMixin):
     command = RoborockCommand.GET_VALLEY_ELECTRICITY_TIMER
     requires_feature = "is_supported_valley_electricity"
 
+    @property
+    def is_on(self) -> bool:
+        """Return whether the Valley Electricity Timer is enabled."""
+        return self.enabled == 1
+
     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_dict())

tests/devices/traits/v1/test_dnd.py

@@ -46,6 +46,7 @@ async def test_get_dnd_timer_success(
     assert dnd_trait.end_hour == 8
     assert dnd_trait.end_minute == 0
     assert dnd_trait.enabled == 1
+    assert dnd_trait.is_on
 
     # Verify the RPC call was made correctly
     mock_rpc_channel.send_command.assert_called_once_with(RoborockCommand.GET_DND_TIMER)
@@ -65,6 +66,7 @@ async def test_get_dnd_timer_disabled(dnd_trait: DoNotDisturbTrait, mock_rpc_cha
     await dnd_trait.refresh()
 
     assert dnd_trait.enabled == 0
+    assert not dnd_trait.is_on
     mock_rpc_channel.send_command.assert_called_once_with(RoborockCommand.GET_DND_TIMER)