Cleanup litterrobot switch entity (#136199)

homeassistant/components/litterrobot/switch.py

@@ -17,18 +17,13 @@ from . import LitterRobotConfigEntry
 from .entity import LitterRobotEntity, _RobotT
 
 
-@dataclass(frozen=True)
+@dataclass(frozen=True, kw_only=True)
-class RequiredKeysMixin(Generic[_RobotT]):
+class RobotSwitchEntityDescription(SwitchEntityDescription, Generic[_RobotT]):
-    """A class that describes robot switch entity required keys."""
-
-    set_fn: Callable[[_RobotT, bool], Coroutine[Any, Any, bool]]
-
-
-@dataclass(frozen=True)
-class RobotSwitchEntityDescription(SwitchEntityDescription, RequiredKeysMixin[_RobotT]):
     """A class that describes robot switch entities."""
 
     entity_category: EntityCategory = EntityCategory.CONFIG
+    set_fn: Callable[[_RobotT, bool], Coroutine[Any, Any, bool]]
+    value_fn: Callable[[_RobotT], bool]
 
 
 ROBOT_SWITCHES = [
@@ -36,34 +31,17 @@ ROBOT_SWITCHES = [
         key="night_light_mode_enabled",
         translation_key="night_light_mode",
         set_fn=lambda robot, value: robot.set_night_light(value),
+        value_fn=lambda robot: robot.night_light_mode_enabled,
     ),
     RobotSwitchEntityDescription[LitterRobot | FeederRobot](
         key="panel_lock_enabled",
         translation_key="panel_lockout",
         set_fn=lambda robot, value: robot.set_panel_lockout(value),
+        value_fn=lambda robot: robot.panel_lock_enabled,
     ),
 ]
 
 
-class RobotSwitchEntity(LitterRobotEntity[_RobotT], SwitchEntity):
-    """Litter-Robot switch entity."""
-
-    entity_description: RobotSwitchEntityDescription[_RobotT]
-
-    @property
-    def is_on(self) -> bool | None:
-        """Return true if switch is on."""
-        return bool(getattr(self.robot, self.entity_description.key))
-
-    async def async_turn_on(self, **kwargs: Any) -> None:
-        """Turn the switch on."""
-        await self.entity_description.set_fn(self.robot, True)
-
-    async def async_turn_off(self, **kwargs: Any) -> None:
-        """Turn the switch off."""
-        await self.entity_description.set_fn(self.robot, False)
-
-
 async def async_setup_entry(
     hass: HomeAssistant,
     entry: LitterRobotConfigEntry,
@@ -78,3 +56,22 @@ async def async_setup_entry(
         if isinstance(robot, (LitterRobot, FeederRobot))
     ]
     async_add_entities(entities)
+
+
+class RobotSwitchEntity(LitterRobotEntity[_RobotT], SwitchEntity):
+    """Litter-Robot switch entity."""
+
+    entity_description: RobotSwitchEntityDescription[_RobotT]
+
+    @property
+    def is_on(self) -> bool | None:
+        """Return true if switch is on."""
+        return self.entity_description.value_fn(self.robot)
+
+    async def async_turn_on(self, **kwargs: Any) -> None:
+        """Turn the switch on."""
+        await self.entity_description.set_fn(self.robot, True)
+
+    async def async_turn_off(self, **kwargs: Any) -> None:
+        """Turn the switch off."""
+        await self.entity_description.set_fn(self.robot, False)