Add pets to litterrobot integration (#136865)

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
+        for _id in itertools.chain(
-        if robot.serial == identifier[1]
+            (robot.serial for robot in entry.runtime_data.account.robots),
+            (pet.id for pet in entry.runtime_data.account.pets),
+        )
+        if _id == identifier[1]
     )

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:

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."""

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

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}"
+        _id = robot.serial if isinstance(robot, Robot) else robot.id
-        self._attr_device_info = DeviceInfo(
+        self._attr_unique_id = f"{_id}-{description.key}"
-            identifiers={(DOMAIN, robot.serial)},
+        self._attr_device_info = get_device_info(robot)
-            manufacturer="Whisker",
-            model=robot.model,
-            name=robot.name,
-            serial_number=robot.serial,
-            sw_version=getattr(robot, "firmware", None),
-        )
 
     async def async_added_to_hass(self) -> None:
         """Set up a listener for the entity."""

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]
+    current_fn: Callable[[_WhiskerEntityT], _CastTypeT | None]
-    options_fn: Callable[[_RobotT], list[_CastTypeT]]
+    options_fn: Callable[[_WhiskerEntityT], list[_CastTypeT]]
-    select_fn: Callable[[_RobotT, str], Coroutine[Any, Any, bool]]
+    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)

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:

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]]
+    set_fn: Callable[[_WhiskerEntityT, bool], Coroutine[Any, Any, bool]]
-    value_fn: Callable[[_RobotT], 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:

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]
+    value_fn: Callable[[_WhiskerEntityT], time | None]
-    set_fn: Callable[[_RobotT, time], Coroutine[Any, Any, bool]]
+    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:

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"

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."""

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