Add set_wait_time command support to Litter-Robot (#48300)

Co-authored-by: J. Nick Koston <nick@koston.org>

homeassistant/components/litterrobot/sensor.py

@@ -1,13 +1,19 @@
 """Support for Litter-Robot sensors."""
 from __future__ import annotations
 
+from typing import Callable
+
 from pylitterbot.robot import Robot
 
 from homeassistant.components.sensor import SensorEntity
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import DEVICE_CLASS_TIMESTAMP, PERCENTAGE
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity import Entity
 
 from .const import DOMAIN
-from .hub import LitterRobotEntity, LitterRobotHub
+from .entity import LitterRobotEntity
+from .hub import LitterRobotHub
 
 
 def icon_for_gauge_level(gauge_level: int | None = None, offset: int = 0) -> str:
@@ -22,66 +28,76 @@ def icon_for_gauge_level(gauge_level: int | None = None, offset: int = 0) -> str
 
 
 class LitterRobotPropertySensor(LitterRobotEntity, SensorEntity):
-    """Litter-Robot property sensors."""
+    """Litter-Robot property sensor."""
 
     def __init__(
         self, robot: Robot, entity_type: str, hub: LitterRobotHub, sensor_attribute: str
-    ):
-        """Pass coordinator to CoordinatorEntity."""
+    ) -> None:
+        """Pass robot, entity_type and hub to LitterRobotEntity."""
         super().__init__(robot, entity_type, hub)
         self.sensor_attribute = sensor_attribute
 
     @property
-    def state(self):
+    def state(self) -> str:
         """Return the state."""
         return getattr(self.robot, self.sensor_attribute)
 
 
 class LitterRobotWasteSensor(LitterRobotPropertySensor):
-    """Litter-Robot sensors."""
+    """Litter-Robot waste sensor."""
 
     @property
-    def unit_of_measurement(self):
+    def unit_of_measurement(self) -> str:
         """Return unit of measurement."""
         return PERCENTAGE
 
     @property
-    def icon(self):
+    def icon(self) -> str:
         """Return the icon to use in the frontend, if any."""
         return icon_for_gauge_level(self.state, 10)
 
 
 class LitterRobotSleepTimeSensor(LitterRobotPropertySensor):
-    """Litter-Robot sleep time sensors."""
+    """Litter-Robot sleep time sensor."""
 
     @property
-    def state(self):
+    def state(self) -> str | None:
         """Return the state."""
-        if self.robot.sleep_mode_active:
+        if self.robot.sleep_mode_enabled:
             return super().state.isoformat()
         return None
 
     @property
-    def device_class(self):
+    def device_class(self) -> str:
         """Return the device class, if any."""
         return DEVICE_CLASS_TIMESTAMP
 
 
-ROBOT_SENSORS = [
-    (LitterRobotWasteSensor, "Waste Drawer", "waste_drawer_gauge"),
+ROBOT_SENSORS: list[tuple[type[LitterRobotPropertySensor], str, str]] = [
+    (LitterRobotWasteSensor, "Waste Drawer", "waste_drawer_level"),
     (LitterRobotSleepTimeSensor, "Sleep Mode Start Time", "sleep_mode_start_time"),
     (LitterRobotSleepTimeSensor, "Sleep Mode End Time", "sleep_mode_end_time"),
 ]
 
 
-async def async_setup_entry(hass, config_entry, async_add_entities):
+async def async_setup_entry(
+    hass: HomeAssistant,
+    entry: ConfigEntry,
+    async_add_entities: Callable[[list[Entity], bool], None],
+) -> None:
     """Set up Litter-Robot sensors using config entry."""
-    hub = hass.data[DOMAIN][config_entry.entry_id]
+    hub: LitterRobotHub = hass.data[DOMAIN][entry.entry_id]
 
     entities = []
     for robot in hub.account.robots:
         for (sensor_class, entity_type, sensor_attribute) in ROBOT_SENSORS:
-            entities.append(sensor_class(robot, entity_type, hub, sensor_attribute))
+            entities.append(
+                sensor_class(
+                    robot=robot,
+                    entity_type=entity_type,
+                    hub=hub,
+                    sensor_attribute=sensor_attribute,
+                )
+            )
 
-    if entities:
-        async_add_entities(entities, True)
+    async_add_entities(entities, True)