Refactor litterrobot to use SensorEntityDescription (#71224)

homeassistant/components/litterrobot/sensor.py

@@ -1,11 +1,19 @@
 """Support for Litter-Robot sensors."""
 from __future__ import annotations
 
+from collections.abc import Callable
+from dataclasses import dataclass
 from datetime import datetime
+from typing import Any
 
 from pylitterbot.robot import Robot
 
-from homeassistant.components.sensor import SensorDeviceClass, SensorEntity, StateType
+from homeassistant.components.sensor import (
+    SensorDeviceClass,
+    SensorEntity,
+    SensorEntityDescription,
+    StateType,
+)
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import PERCENTAGE
 from homeassistant.core import HomeAssistant
@@ -27,56 +35,64 @@ def icon_for_gauge_level(gauge_level: int | None = None, offset: int = 0) -> str
     return "mdi:gauge-low"
 
 
-class LitterRobotPropertySensor(LitterRobotEntity, SensorEntity):
+@dataclass
-    """Litter-Robot property sensor."""
+class LitterRobotSensorEntityDescription(SensorEntityDescription):
+    """A class that describes Litter-Robot sensor entities."""
+
+    icon_fn: Callable[[Any], str | None] = lambda _: None
+    should_report: Callable[[Robot], bool] = lambda _: True
+
+
+class LitterRobotSensorEntity(LitterRobotEntity, SensorEntity):
+    """Litter-Robot sensor entity."""
+
+    entity_description: LitterRobotSensorEntityDescription
 
     def __init__(
-        self, robot: Robot, entity_type: str, hub: LitterRobotHub, sensor_attribute: str
+        self,
+        robot: Robot,
+        hub: LitterRobotHub,
+        description: LitterRobotSensorEntityDescription,
     ) -> None:
-        """Pass robot, entity_type and hub to LitterRobotEntity."""
+        """Initialize a Litter-Robot sensor entity."""
-        super().__init__(robot, entity_type, hub)
+        assert description.name
-        self.sensor_attribute = sensor_attribute
+        super().__init__(robot, description.name, hub)
+        self.entity_description = description
 
     @property
     def native_value(self) -> StateType | datetime:
         """Return the state."""
-        return getattr(self.robot, self.sensor_attribute)
+        if self.entity_description.should_report(self.robot):
-
+            return getattr(self.robot, self.entity_description.key)
-
-class LitterRobotWasteSensor(LitterRobotPropertySensor):
-    """Litter-Robot waste sensor."""
-
-    @property
-    def native_unit_of_measurement(self) -> str:
-        """Return unit of measurement."""
-        return PERCENTAGE
-
-    @property
-    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 sensor."""
-
-    @property
-    def native_value(self) -> StateType | datetime:
-        """Return the state."""
-        if self.robot.sleep_mode_enabled:
-            return super().native_value
         return None
 
     @property
-    def device_class(self) -> str:
+    def icon(self) -> str | None:
-        """Return the device class, if any."""
+        """Return the icon to use in the frontend, if any."""
-        return SensorDeviceClass.TIMESTAMP
+        if (icon := self.entity_description.icon_fn(self.state)) is not None:
+            return icon
+        return super().icon
 
 
-ROBOT_SENSORS: list[tuple[type[LitterRobotPropertySensor], str, str]] = [
+ROBOT_SENSORS = [
-    (LitterRobotWasteSensor, "Waste Drawer", "waste_drawer_level"),
+    LitterRobotSensorEntityDescription(
-    (LitterRobotSleepTimeSensor, "Sleep Mode Start Time", "sleep_mode_start_time"),
+        name="Waste Drawer",
-    (LitterRobotSleepTimeSensor, "Sleep Mode End Time", "sleep_mode_end_time"),
+        key="waste_drawer_level",
+        native_unit_of_measurement=PERCENTAGE,
+        icon_fn=lambda state: icon_for_gauge_level(state, 10),
+    ),
+    LitterRobotSensorEntityDescription(
+        name="Sleep Mode Start Time",
+        key="sleep_mode_start_time",
+        device_class=SensorDeviceClass.TIMESTAMP,
+        should_report=lambda robot: robot.sleep_mode_enabled,
+    ),
+    LitterRobotSensorEntityDescription(
+        name="Sleep Mode End Time",
+        key="sleep_mode_end_time",
+        device_class=SensorDeviceClass.TIMESTAMP,
+        should_report=lambda robot: robot.sleep_mode_enabled,
+    ),
 ]
 
 
@@ -87,17 +103,8 @@ async def async_setup_entry(
 ) -> None:
     """Set up Litter-Robot sensors using config entry."""
     hub: LitterRobotHub = hass.data[DOMAIN][entry.entry_id]
-
+    async_add_entities(
-    entities = []
+        LitterRobotSensorEntity(robot=robot, hub=hub, description=description)
-    for robot in hub.account.robots:
+        for description in ROBOT_SENSORS
-        for (sensor_class, entity_type, sensor_attribute) in ROBOT_SENSORS:
+        for robot in hub.account.robots
-            entities.append(
+    )
-                sensor_class(
-                    robot=robot,
-                    entity_type=entity_type,
-                    hub=hub,
-                    sensor_attribute=sensor_attribute,
-                )
-            )
-
-    async_add_entities(entities)

tests/components/litterrobot/conftest.py

@@ -67,6 +67,12 @@ def mock_account_with_sleeping_robot() -> MagicMock:
     return create_mock_account({"sleepModeActive": "102:00:00"})
 
 
+@pytest.fixture
+def mock_account_with_sleep_disabled_robot() -> MagicMock:
+    """Mock a Litter-Robot account with a robot that has sleep mode disabled."""
+    return create_mock_account({"sleepModeActive": "0"})
+
+
 @pytest.fixture
 def mock_account_with_robot_not_recently_seen() -> MagicMock:
     """Mock a Litter-Robot account with a sleeping robot."""

tests/components/litterrobot/test_sensor.py

@@ -1,16 +1,19 @@
 """Test the Litter-Robot sensor entity."""
-from unittest.mock import Mock
+from unittest.mock import MagicMock
 
-from homeassistant.components.litterrobot.sensor import LitterRobotSleepTimeSensor
 from homeassistant.components.sensor import DOMAIN as PLATFORM_DOMAIN, SensorDeviceClass
-from homeassistant.const import PERCENTAGE
+from homeassistant.const import PERCENTAGE, STATE_UNKNOWN
+from homeassistant.core import HomeAssistant
 
-from .conftest import create_mock_robot, setup_integration
+from .conftest import setup_integration
 
 WASTE_DRAWER_ENTITY_ID = "sensor.test_waste_drawer"
+SLEEP_START_TIME_ENTITY_ID = "sensor.test_sleep_mode_start_time"
 
 
-async def test_waste_drawer_sensor(hass, mock_account):
+async def test_waste_drawer_sensor(
+    hass: HomeAssistant, mock_account: MagicMock
+) -> None:
     """Tests the waste drawer sensor entity was set up."""
     await setup_integration(hass, mock_account, PLATFORM_DOMAIN)
 
@@ -20,20 +23,21 @@ async def test_waste_drawer_sensor(hass, mock_account):
     assert sensor.attributes["unit_of_measurement"] == PERCENTAGE
 
 
-async def test_sleep_time_sensor_with_none_state(hass):
+async def test_sleep_time_sensor_with_sleep_disabled(
-    """Tests the sleep mode start time sensor where sleep mode is inactive."""
+    hass: HomeAssistant, mock_account_with_sleep_disabled_robot: MagicMock
-    robot = create_mock_robot({"sleepModeActive": "0"})
+) -> None:
-    sensor = LitterRobotSleepTimeSensor(
+    """Tests the sleep mode start time sensor where sleep mode is disabled."""
-        robot, "Sleep Mode Start Time", Mock(), "sleep_mode_start_time"
+    await setup_integration(
+        hass, mock_account_with_sleep_disabled_robot, PLATFORM_DOMAIN
     )
-    sensor.hass = hass
 
+    sensor = hass.states.get(SLEEP_START_TIME_ENTITY_ID)
     assert sensor
-    assert sensor.state is None
+    assert sensor.state == STATE_UNKNOWN
-    assert sensor.device_class is SensorDeviceClass.TIMESTAMP
+    assert sensor.attributes["device_class"] == SensorDeviceClass.TIMESTAMP
 
 
-async def test_gauge_icon():
+async def test_gauge_icon() -> None:
     """Test icon generator for gauge sensor."""
     from homeassistant.components.litterrobot.sensor import icon_for_gauge_level