Add pets to litterrobot integration (#136865)

pull/137022/head
Nathan Spencer 2025-01-31 09:35:08 -07:00 committed by GitHub
parent e18dc063ba
commit b1c3d0857a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
12 changed files with 133 additions and 55 deletions

View File

@ -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]
)

View File

@ -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:

View File

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

View File

@ -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

View File

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

View File

@ -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)

View File

@ -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:

View File

@ -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:

View File

@ -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:

View File

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

View File

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

View File

@ -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