From b1c3d0857a712f6f835c3de07169c4a0db693b21 Mon Sep 17 00:00:00 2001 From: Nathan Spencer Date: Fri, 31 Jan 2025 09:35:08 -0700 Subject: [PATCH] Add pets to litterrobot integration (#136865) --- .../components/litterrobot/__init__.py | 9 ++++- .../components/litterrobot/binary_sensor.py | 12 +++--- .../components/litterrobot/button.py | 10 ++--- .../components/litterrobot/coordinator.py | 2 + .../components/litterrobot/entity.py | 40 +++++++++++++------ .../components/litterrobot/select.py | 20 +++++----- .../components/litterrobot/sensor.py | 32 +++++++++++---- .../components/litterrobot/switch.py | 12 +++--- homeassistant/components/litterrobot/time.py | 12 +++--- tests/components/litterrobot/common.py | 10 +++++ tests/components/litterrobot/conftest.py | 19 ++++++++- tests/components/litterrobot/test_sensor.py | 10 +++++ 12 files changed, 133 insertions(+), 55 deletions(-) diff --git a/homeassistant/components/litterrobot/__init__.py b/homeassistant/components/litterrobot/__init__.py index 1f926d37a61..2823450d9ad 100644 --- a/homeassistant/components/litterrobot/__init__.py +++ b/homeassistant/components/litterrobot/__init__.py @@ -2,6 +2,8 @@ from __future__ import annotations +import itertools + from homeassistant.const import Platform from homeassistant.core import HomeAssistant from homeassistant.helpers.device_registry import DeviceEntry @@ -46,6 +48,9 @@ async def async_remove_config_entry_device( identifier for identifier in device_entry.identifiers if identifier[0] == DOMAIN - for robot in entry.runtime_data.account.robots - if robot.serial == identifier[1] + for _id in itertools.chain( + (robot.serial for robot in entry.runtime_data.account.robots), + (pet.id for pet in entry.runtime_data.account.pets), + ) + if _id == identifier[1] ) diff --git a/homeassistant/components/litterrobot/binary_sensor.py b/homeassistant/components/litterrobot/binary_sensor.py index e6cf23fa27c..700985d285f 100644 --- a/homeassistant/components/litterrobot/binary_sensor.py +++ b/homeassistant/components/litterrobot/binary_sensor.py @@ -18,16 +18,16 @@ from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback from .coordinator import LitterRobotConfigEntry -from .entity import LitterRobotEntity, _RobotT +from .entity import LitterRobotEntity, _WhiskerEntityT @dataclass(frozen=True, kw_only=True) class RobotBinarySensorEntityDescription( - BinarySensorEntityDescription, Generic[_RobotT] + BinarySensorEntityDescription, Generic[_WhiskerEntityT] ): """A class that describes robot binary sensor entities.""" - is_on_fn: Callable[[_RobotT], bool] + is_on_fn: Callable[[_WhiskerEntityT], bool] BINARY_SENSOR_MAP: dict[type[Robot], tuple[RobotBinarySensorEntityDescription, ...]] = { @@ -78,10 +78,12 @@ async def async_setup_entry( ) -class LitterRobotBinarySensorEntity(LitterRobotEntity[_RobotT], BinarySensorEntity): +class LitterRobotBinarySensorEntity( + LitterRobotEntity[_WhiskerEntityT], BinarySensorEntity +): """Litter-Robot binary sensor entity.""" - entity_description: RobotBinarySensorEntityDescription[_RobotT] + entity_description: RobotBinarySensorEntityDescription[_WhiskerEntityT] @property def is_on(self) -> bool: diff --git a/homeassistant/components/litterrobot/button.py b/homeassistant/components/litterrobot/button.py index 01888e7fbae..758548b3a67 100644 --- a/homeassistant/components/litterrobot/button.py +++ b/homeassistant/components/litterrobot/button.py @@ -14,14 +14,14 @@ from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback from .coordinator import LitterRobotConfigEntry -from .entity import LitterRobotEntity, _RobotT +from .entity import LitterRobotEntity, _WhiskerEntityT @dataclass(frozen=True, kw_only=True) -class RobotButtonEntityDescription(ButtonEntityDescription, Generic[_RobotT]): +class RobotButtonEntityDescription(ButtonEntityDescription, Generic[_WhiskerEntityT]): """A class that describes robot button entities.""" - press_fn: Callable[[_RobotT], Coroutine[Any, Any, bool]] + press_fn: Callable[[_WhiskerEntityT], Coroutine[Any, Any, bool]] ROBOT_BUTTON_MAP: dict[type[Robot], RobotButtonEntityDescription] = { @@ -62,10 +62,10 @@ async def async_setup_entry( ) -class LitterRobotButtonEntity(LitterRobotEntity[_RobotT], ButtonEntity): +class LitterRobotButtonEntity(LitterRobotEntity[_WhiskerEntityT], ButtonEntity): """Litter-Robot button entity.""" - entity_description: RobotButtonEntityDescription[_RobotT] + entity_description: RobotButtonEntityDescription[_WhiskerEntityT] async def async_press(self) -> None: """Press the button.""" diff --git a/homeassistant/components/litterrobot/coordinator.py b/homeassistant/components/litterrobot/coordinator.py index a56a6607d32..c99d4794ff6 100644 --- a/homeassistant/components/litterrobot/coordinator.py +++ b/homeassistant/components/litterrobot/coordinator.py @@ -47,6 +47,7 @@ class LitterRobotDataUpdateCoordinator(DataUpdateCoordinator[None]): async def _async_update_data(self) -> None: """Update all device states from the Litter-Robot API.""" await self.account.refresh_robots() + await self.account.load_pets() async def _async_setup(self) -> None: """Set up the coordinator.""" @@ -56,6 +57,7 @@ class LitterRobotDataUpdateCoordinator(DataUpdateCoordinator[None]): password=self.config_entry.data[CONF_PASSWORD], load_robots=True, subscribe_for_updates=True, + load_pets=True, ) except LitterRobotLoginException as ex: raise ConfigEntryAuthFailed("Invalid credentials") from ex diff --git a/homeassistant/components/litterrobot/entity.py b/homeassistant/components/litterrobot/entity.py index 36cbbb730ce..9e9cc8f0740 100644 --- a/homeassistant/components/litterrobot/entity.py +++ b/homeassistant/components/litterrobot/entity.py @@ -4,7 +4,7 @@ from __future__ import annotations from typing import Generic, TypeVar -from pylitterbot import Robot +from pylitterbot import Pet, Robot from pylitterbot.robot import EVENT_UPDATE from homeassistant.helpers.device_registry import DeviceInfo @@ -14,11 +14,31 @@ from homeassistant.helpers.update_coordinator import CoordinatorEntity from .const import DOMAIN from .coordinator import LitterRobotDataUpdateCoordinator -_RobotT = TypeVar("_RobotT", bound=Robot) +_WhiskerEntityT = TypeVar("_WhiskerEntityT", bound=Robot | Pet) + + +def get_device_info(whisker_entity: _WhiskerEntityT) -> DeviceInfo: + """Get device info for a robot or pet.""" + if isinstance(whisker_entity, Robot): + return DeviceInfo( + identifiers={(DOMAIN, whisker_entity.serial)}, + manufacturer="Whisker", + model=whisker_entity.model, + name=whisker_entity.name, + serial_number=whisker_entity.serial, + sw_version=getattr(whisker_entity, "firmware", None), + ) + breed = ", ".join(breed for breed in whisker_entity.breeds or []) + return DeviceInfo( + identifiers={(DOMAIN, whisker_entity.id)}, + manufacturer="Whisker", + model=f"{breed} {whisker_entity.pet_type}".strip().capitalize(), + name=whisker_entity.name, + ) class LitterRobotEntity( - CoordinatorEntity[LitterRobotDataUpdateCoordinator], Generic[_RobotT] + CoordinatorEntity[LitterRobotDataUpdateCoordinator], Generic[_WhiskerEntityT] ): """Generic Litter-Robot entity representing common data and methods.""" @@ -26,7 +46,7 @@ class LitterRobotEntity( def __init__( self, - robot: _RobotT, + robot: _WhiskerEntityT, coordinator: LitterRobotDataUpdateCoordinator, description: EntityDescription, ) -> None: @@ -34,15 +54,9 @@ class LitterRobotEntity( super().__init__(coordinator) self.robot = robot self.entity_description = description - self._attr_unique_id = f"{robot.serial}-{description.key}" - self._attr_device_info = DeviceInfo( - identifiers={(DOMAIN, robot.serial)}, - manufacturer="Whisker", - model=robot.model, - name=robot.name, - serial_number=robot.serial, - sw_version=getattr(robot, "firmware", None), - ) + _id = robot.serial if isinstance(robot, Robot) else robot.id + self._attr_unique_id = f"{_id}-{description.key}" + self._attr_device_info = get_device_info(robot) async def async_added_to_hass(self) -> None: """Set up a listener for the entity.""" diff --git a/homeassistant/components/litterrobot/select.py b/homeassistant/components/litterrobot/select.py index 1a3d2fc2fb4..f6e3781f3df 100644 --- a/homeassistant/components/litterrobot/select.py +++ b/homeassistant/components/litterrobot/select.py @@ -15,21 +15,21 @@ from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback from .coordinator import LitterRobotConfigEntry, LitterRobotDataUpdateCoordinator -from .entity import LitterRobotEntity, _RobotT +from .entity import LitterRobotEntity, _WhiskerEntityT _CastTypeT = TypeVar("_CastTypeT", int, float, str) @dataclass(frozen=True, kw_only=True) class RobotSelectEntityDescription( - SelectEntityDescription, Generic[_RobotT, _CastTypeT] + SelectEntityDescription, Generic[_WhiskerEntityT, _CastTypeT] ): """A class that describes robot select entities.""" entity_category: EntityCategory = EntityCategory.CONFIG - current_fn: Callable[[_RobotT], _CastTypeT | None] - options_fn: Callable[[_RobotT], list[_CastTypeT]] - select_fn: Callable[[_RobotT, str], Coroutine[Any, Any, bool]] + current_fn: Callable[[_WhiskerEntityT], _CastTypeT | None] + options_fn: Callable[[_WhiskerEntityT], list[_CastTypeT]] + select_fn: Callable[[_WhiskerEntityT, str], Coroutine[Any, Any, bool]] ROBOT_SELECT_MAP: dict[type[Robot], RobotSelectEntityDescription] = { @@ -83,17 +83,19 @@ async def async_setup_entry( class LitterRobotSelectEntity( - LitterRobotEntity[_RobotT], SelectEntity, Generic[_RobotT, _CastTypeT] + LitterRobotEntity[_WhiskerEntityT], + SelectEntity, + Generic[_WhiskerEntityT, _CastTypeT], ): """Litter-Robot Select.""" - entity_description: RobotSelectEntityDescription[_RobotT, _CastTypeT] + entity_description: RobotSelectEntityDescription[_WhiskerEntityT, _CastTypeT] def __init__( self, - robot: _RobotT, + robot: _WhiskerEntityT, coordinator: LitterRobotDataUpdateCoordinator, - description: RobotSelectEntityDescription[_RobotT, _CastTypeT], + description: RobotSelectEntityDescription[_WhiskerEntityT, _CastTypeT], ) -> None: """Initialize a Litter-Robot select entity.""" super().__init__(robot, coordinator, description) diff --git a/homeassistant/components/litterrobot/sensor.py b/homeassistant/components/litterrobot/sensor.py index 6545d7c7ae7..3e25a0556c6 100644 --- a/homeassistant/components/litterrobot/sensor.py +++ b/homeassistant/components/litterrobot/sensor.py @@ -7,7 +7,7 @@ from dataclasses import dataclass from datetime import datetime from typing import Any, Generic -from pylitterbot import FeederRobot, LitterRobot, LitterRobot4, Robot +from pylitterbot import FeederRobot, LitterRobot, LitterRobot4, Pet, Robot from homeassistant.components.sensor import ( SensorDeviceClass, @@ -20,7 +20,7 @@ from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback from .coordinator import LitterRobotConfigEntry -from .entity import LitterRobotEntity, _RobotT +from .entity import LitterRobotEntity, _WhiskerEntityT def icon_for_gauge_level(gauge_level: int | None = None, offset: int = 0) -> str: @@ -35,11 +35,11 @@ def icon_for_gauge_level(gauge_level: int | None = None, offset: int = 0) -> str @dataclass(frozen=True, kw_only=True) -class RobotSensorEntityDescription(SensorEntityDescription, Generic[_RobotT]): +class RobotSensorEntityDescription(SensorEntityDescription, Generic[_WhiskerEntityT]): """A class that describes robot sensor entities.""" icon_fn: Callable[[Any], str | None] = lambda _: None - value_fn: Callable[[_RobotT], float | datetime | str | None] + value_fn: Callable[[_WhiskerEntityT], float | datetime | str | None] ROBOT_SENSOR_MAP: dict[type[Robot], list[RobotSensorEntityDescription]] = { @@ -146,6 +146,16 @@ ROBOT_SENSOR_MAP: dict[type[Robot], list[RobotSensorEntityDescription]] = { ], } +PET_SENSORS: list[RobotSensorEntityDescription] = [ + RobotSensorEntityDescription[Pet]( + key="weight", + device_class=SensorDeviceClass.WEIGHT, + native_unit_of_measurement=UnitOfMass.POUNDS, + state_class=SensorStateClass.MEASUREMENT, + value_fn=lambda pet: pet.weight, + ) +] + async def async_setup_entry( hass: HomeAssistant, @@ -154,7 +164,7 @@ async def async_setup_entry( ) -> None: """Set up Litter-Robot sensors using config entry.""" coordinator = entry.runtime_data - async_add_entities( + entities: list[LitterRobotSensorEntity] = [ LitterRobotSensorEntity( robot=robot, coordinator=coordinator, description=description ) @@ -162,13 +172,21 @@ async def async_setup_entry( for robot_type, entity_descriptions in ROBOT_SENSOR_MAP.items() if isinstance(robot, robot_type) for description in entity_descriptions + ] + entities.extend( + LitterRobotSensorEntity( + robot=pet, coordinator=coordinator, description=description + ) + for pet in coordinator.account.pets + for description in PET_SENSORS ) + async_add_entities(entities) -class LitterRobotSensorEntity(LitterRobotEntity[_RobotT], SensorEntity): +class LitterRobotSensorEntity(LitterRobotEntity[_WhiskerEntityT], SensorEntity): """Litter-Robot sensor entity.""" - entity_description: RobotSensorEntityDescription[_RobotT] + entity_description: RobotSensorEntityDescription[_WhiskerEntityT] @property def native_value(self) -> float | datetime | str | None: diff --git a/homeassistant/components/litterrobot/switch.py b/homeassistant/components/litterrobot/switch.py index 7ded89d552b..4839748c068 100644 --- a/homeassistant/components/litterrobot/switch.py +++ b/homeassistant/components/litterrobot/switch.py @@ -14,16 +14,16 @@ from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback from .coordinator import LitterRobotConfigEntry -from .entity import LitterRobotEntity, _RobotT +from .entity import LitterRobotEntity, _WhiskerEntityT @dataclass(frozen=True, kw_only=True) -class RobotSwitchEntityDescription(SwitchEntityDescription, Generic[_RobotT]): +class RobotSwitchEntityDescription(SwitchEntityDescription, Generic[_WhiskerEntityT]): """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] + set_fn: Callable[[_WhiskerEntityT, bool], Coroutine[Any, Any, bool]] + value_fn: Callable[[_WhiskerEntityT], bool] ROBOT_SWITCHES = [ @@ -57,10 +57,10 @@ async def async_setup_entry( ) -class RobotSwitchEntity(LitterRobotEntity[_RobotT], SwitchEntity): +class RobotSwitchEntity(LitterRobotEntity[_WhiskerEntityT], SwitchEntity): """Litter-Robot switch entity.""" - entity_description: RobotSwitchEntityDescription[_RobotT] + entity_description: RobotSwitchEntityDescription[_WhiskerEntityT] @property def is_on(self) -> bool | None: diff --git a/homeassistant/components/litterrobot/time.py b/homeassistant/components/litterrobot/time.py index 3fa93b14dd9..69d81d63eae 100644 --- a/homeassistant/components/litterrobot/time.py +++ b/homeassistant/components/litterrobot/time.py @@ -16,15 +16,15 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.util import dt as dt_util from .coordinator import LitterRobotConfigEntry -from .entity import LitterRobotEntity, _RobotT +from .entity import LitterRobotEntity, _WhiskerEntityT @dataclass(frozen=True, kw_only=True) -class RobotTimeEntityDescription(TimeEntityDescription, Generic[_RobotT]): +class RobotTimeEntityDescription(TimeEntityDescription, Generic[_WhiskerEntityT]): """A class that describes robot time entities.""" - value_fn: Callable[[_RobotT], time | None] - set_fn: Callable[[_RobotT, time], Coroutine[Any, Any, bool]] + value_fn: Callable[[_WhiskerEntityT], time | None] + set_fn: Callable[[_WhiskerEntityT, time], Coroutine[Any, Any, bool]] def _as_local_time(start: datetime | None) -> time | None: @@ -64,10 +64,10 @@ async def async_setup_entry( ) -class LitterRobotTimeEntity(LitterRobotEntity[_RobotT], TimeEntity): +class LitterRobotTimeEntity(LitterRobotEntity[_WhiskerEntityT], TimeEntity): """Litter-Robot time entity.""" - entity_description: RobotTimeEntityDescription[_RobotT] + entity_description: RobotTimeEntityDescription[_WhiskerEntityT] @property def native_value(self) -> time | None: diff --git a/tests/components/litterrobot/common.py b/tests/components/litterrobot/common.py index b29fa753801..d96ce06ca59 100644 --- a/tests/components/litterrobot/common.py +++ b/tests/components/litterrobot/common.py @@ -150,5 +150,15 @@ FEEDER_ROBOT_DATA = { }, ], } +PET_DATA = { + "petId": "PET-123", + "userId": "1234567", + "createdAt": "2023-04-27T23:26:49.813Z", + "name": "Kitty", + "type": "CAT", + "gender": "FEMALE", + "lastWeightReading": 9.1, + "breeds": ["sphynx"], +} VACUUM_ENTITY_ID = "vacuum.test_litter_box" diff --git a/tests/components/litterrobot/conftest.py b/tests/components/litterrobot/conftest.py index e60e0cbd36d..d22c4b2ec49 100644 --- a/tests/components/litterrobot/conftest.py +++ b/tests/components/litterrobot/conftest.py @@ -5,13 +5,20 @@ from __future__ import annotations from typing import Any from unittest.mock import AsyncMock, MagicMock, patch -from pylitterbot import Account, FeederRobot, LitterRobot3, LitterRobot4, Robot +from pylitterbot import Account, FeederRobot, LitterRobot3, LitterRobot4, Pet, Robot from pylitterbot.exceptions import InvalidCommandException import pytest from homeassistant.core import HomeAssistant -from .common import CONFIG, DOMAIN, FEEDER_ROBOT_DATA, ROBOT_4_DATA, ROBOT_DATA +from .common import ( + CONFIG, + DOMAIN, + FEEDER_ROBOT_DATA, + PET_DATA, + ROBOT_4_DATA, + ROBOT_DATA, +) from tests.common import MockConfigEntry @@ -50,6 +57,7 @@ def create_mock_account( skip_robots: bool = False, v4: bool = False, feeder: bool = False, + pet: bool = False, ) -> MagicMock: """Create a mock Litter-Robot account.""" account = MagicMock(spec=Account) @@ -60,6 +68,7 @@ def create_mock_account( if skip_robots else [create_mock_robot(robot_data, account, v4, feeder, side_effect)] ) + account.pets = [Pet(PET_DATA, account.session)] if pet else [] return account @@ -81,6 +90,12 @@ def mock_account_with_feederrobot() -> MagicMock: return create_mock_account(feeder=True) +@pytest.fixture +def mock_account_with_pet() -> MagicMock: + """Mock account with Feeder-Robot.""" + return create_mock_account(pet=True) + + @pytest.fixture def mock_account_with_no_robots() -> MagicMock: """Mock a Litter-Robot account.""" diff --git a/tests/components/litterrobot/test_sensor.py b/tests/components/litterrobot/test_sensor.py index 360d13096a7..e290d96fcf4 100644 --- a/tests/components/litterrobot/test_sensor.py +++ b/tests/components/litterrobot/test_sensor.py @@ -104,3 +104,13 @@ async def test_feeder_robot_sensor( sensor = hass.states.get("sensor.test_food_level") assert sensor.state == "10" assert sensor.attributes["unit_of_measurement"] == PERCENTAGE + + +async def test_pet_weight_sensor( + hass: HomeAssistant, mock_account_with_pet: MagicMock +) -> None: + """Tests pet weight sensors.""" + await setup_integration(hass, mock_account_with_pet, PLATFORM_DOMAIN) + sensor = hass.states.get("sensor.kitty_weight") + assert sensor.state == "9.1" + assert sensor.attributes["unit_of_measurement"] == UnitOfMass.POUNDS