diff --git a/homeassistant/components/litterrobot/sensor.py b/homeassistant/components/litterrobot/sensor.py index de3126986d1..751d16551c3 100644 --- a/homeassistant/components/litterrobot/sensor.py +++ b/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): - """Litter-Robot property sensor.""" +@dataclass +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.""" - super().__init__(robot, entity_type, hub) - self.sensor_attribute = sensor_attribute + """Initialize a Litter-Robot sensor entity.""" + assert description.name + 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) - - -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 + if self.entity_description.should_report(self.robot): + return getattr(self.robot, self.entity_description.key) return None @property - def device_class(self) -> str: - """Return the device class, if any.""" - return SensorDeviceClass.TIMESTAMP + def icon(self) -> str | None: + """Return the icon to use in the frontend, if any.""" + 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]] = [ - (LitterRobotWasteSensor, "Waste Drawer", "waste_drawer_level"), - (LitterRobotSleepTimeSensor, "Sleep Mode Start Time", "sleep_mode_start_time"), - (LitterRobotSleepTimeSensor, "Sleep Mode End Time", "sleep_mode_end_time"), +ROBOT_SENSORS = [ + LitterRobotSensorEntityDescription( + name="Waste Drawer", + 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] - - entities = [] - for robot in hub.account.robots: - for (sensor_class, entity_type, sensor_attribute) in ROBOT_SENSORS: - entities.append( - sensor_class( - robot=robot, - entity_type=entity_type, - hub=hub, - sensor_attribute=sensor_attribute, - ) - ) - - async_add_entities(entities) + async_add_entities( + LitterRobotSensorEntity(robot=robot, hub=hub, description=description) + for description in ROBOT_SENSORS + for robot in hub.account.robots + ) diff --git a/tests/components/litterrobot/conftest.py b/tests/components/litterrobot/conftest.py index 839ffe2952d..e8ec5324ae6 100644 --- a/tests/components/litterrobot/conftest.py +++ b/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.""" diff --git a/tests/components/litterrobot/test_sensor.py b/tests/components/litterrobot/test_sensor.py index e3e62d5f5e4..ce91541f0d7 100644 --- a/tests/components/litterrobot/test_sensor.py +++ b/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): - """Tests the sleep mode start time sensor where sleep mode is inactive.""" - robot = create_mock_robot({"sleepModeActive": "0"}) - sensor = LitterRobotSleepTimeSensor( - robot, "Sleep Mode Start Time", Mock(), "sleep_mode_start_time" +async def test_sleep_time_sensor_with_sleep_disabled( + hass: HomeAssistant, mock_account_with_sleep_disabled_robot: MagicMock +) -> None: + """Tests the sleep mode start time sensor where sleep mode is disabled.""" + 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.device_class is SensorDeviceClass.TIMESTAMP + assert sensor.state == STATE_UNKNOWN + 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