Code quality improvements for litterrobot integration (#77605)

pull/77679/head
Nathan Spencer 2022-09-01 12:02:46 -06:00 committed by GitHub
parent 1692808d5b
commit d0d1b303fd
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
7 changed files with 123 additions and 94 deletions

View File

@ -1,21 +1,25 @@
"""Support for Litter-Robot button."""
from __future__ import annotations
from collections.abc import Callable, Coroutine, Iterable
from collections.abc import Callable, Coroutine
from dataclasses import dataclass
import itertools
from typing import Any, Generic
from pylitterbot import FeederRobot, LitterRobot3
from homeassistant.components.button import ButtonEntity, ButtonEntityDescription
from homeassistant.components.button import (
DOMAIN as PLATFORM,
ButtonEntity,
ButtonEntityDescription,
)
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity import EntityCategory
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from .const import DOMAIN
from .entity import LitterRobotEntity, _RobotT
from .entity import LitterRobotEntity, _RobotT, async_update_unique_id
from .hub import LitterRobotHub
@ -26,7 +30,8 @@ async def async_setup_entry(
) -> None:
"""Set up Litter-Robot cleaner using config entry."""
hub: LitterRobotHub = hass.data[DOMAIN][entry.entry_id]
entities: Iterable[LitterRobotButtonEntity] = itertools.chain(
entities: list[LitterRobotButtonEntity] = list(
itertools.chain(
(
LitterRobotButtonEntity(
robot=robot, hub=hub, description=LITTER_ROBOT_BUTTON
@ -41,6 +46,8 @@ async def async_setup_entry(
for robot in hub.feeder_robots()
),
)
)
async_update_unique_id(hass, PLATFORM, entities)
async_add_entities(entities)
@ -76,17 +83,6 @@ class LitterRobotButtonEntity(LitterRobotEntity[_RobotT], ButtonEntity):
entity_description: RobotButtonEntityDescription[_RobotT]
def __init__(
self,
robot: _RobotT,
hub: LitterRobotHub,
description: RobotButtonEntityDescription[_RobotT],
) -> None:
"""Initialize a Litter-Robot button entity."""
assert description.name
super().__init__(robot, description.name, hub)
self.entity_description = description
async def async_press(self) -> None:
"""Press the button."""
await self.entity_description.press_fn(self.robot)

View File

@ -1,7 +1,7 @@
"""Litter-Robot entities for common data and methods."""
from __future__ import annotations
from collections.abc import Callable, Coroutine
from collections.abc import Callable, Coroutine, Iterable
from datetime import time
import logging
from typing import Any, Generic, TypeVar
@ -10,8 +10,9 @@ from pylitterbot import Robot
from pylitterbot.exceptions import InvalidCommandException
from typing_extensions import ParamSpec
from homeassistant.core import CALLBACK_TYPE, callback
from homeassistant.helpers.entity import DeviceInfo, EntityCategory
from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback
from homeassistant.helpers.entity import DeviceInfo, EntityCategory, EntityDescription
import homeassistant.helpers.entity_registry as er
from homeassistant.helpers.event import async_call_later
from homeassistant.helpers.update_coordinator import (
CoordinatorEntity,
@ -36,18 +37,18 @@ class LitterRobotEntity(
_attr_has_entity_name = True
def __init__(self, robot: _RobotT, entity_type: str, hub: LitterRobotHub) -> None:
def __init__(
self, robot: _RobotT, hub: LitterRobotHub, description: EntityDescription
) -> None:
"""Pass coordinator to CoordinatorEntity."""
super().__init__(hub.coordinator)
self.robot = robot
self.entity_type = entity_type
self.hub = hub
self._attr_name = entity_type.capitalize()
@property
def unique_id(self) -> str:
"""Return a unique ID."""
return f"{self.robot.serial}-{self.entity_type}"
self.entity_description = description
self._attr_unique_id = f"{self.robot.serial}-{description.key}"
# The following can be removed in 2022.12 after adjusting names in entities appropriately
if description.name is not None:
self._attr_name = description.name.capitalize()
@property
def device_info(self) -> DeviceInfo:
@ -65,9 +66,11 @@ class LitterRobotEntity(
class LitterRobotControlEntity(LitterRobotEntity[_RobotT]):
"""A Litter-Robot entity that can control the unit."""
def __init__(self, robot: _RobotT, entity_type: str, hub: LitterRobotHub) -> None:
def __init__(
self, robot: _RobotT, hub: LitterRobotHub, description: EntityDescription
) -> None:
"""Init a Litter-Robot control entity."""
super().__init__(robot=robot, entity_type=entity_type, hub=hub)
super().__init__(robot=robot, hub=hub, description=description)
self._refresh_callback: CALLBACK_TYPE | None = None
async def perform_action_and_refresh(
@ -134,9 +137,11 @@ class LitterRobotConfigEntity(LitterRobotControlEntity[_RobotT]):
_attr_entity_category = EntityCategory.CONFIG
def __init__(self, robot: _RobotT, entity_type: str, hub: LitterRobotHub) -> None:
def __init__(
self, robot: _RobotT, hub: LitterRobotHub, description: EntityDescription
) -> None:
"""Init a Litter-Robot control entity."""
super().__init__(robot=robot, entity_type=entity_type, hub=hub)
super().__init__(robot=robot, hub=hub, description=description)
self._assumed_state: bool | None = None
async def perform_action_and_assume_state(
@ -146,3 +151,18 @@ class LitterRobotConfigEntity(LitterRobotControlEntity[_RobotT]):
if await self.perform_action_and_refresh(action, assumed_state):
self._assumed_state = assumed_state
self.async_write_ha_state()
def async_update_unique_id(
hass: HomeAssistant, domain: str, entities: Iterable[LitterRobotEntity[_RobotT]]
) -> None:
"""Update unique ID to be based on entity description key instead of name.
Introduced with release 2022.9.
"""
ent_reg = er.async_get(hass)
for entity in entities:
old_unique_id = f"{entity.robot.serial}-{entity.entity_description.name}"
if entity_id := ent_reg.async_get_entity_id(domain, DOMAIN, old_unique_id):
new_unique_id = f"{entity.robot.serial}-{entity.entity_description.key}"
ent_reg.async_update_entity(entity_id, new_unique_id=new_unique_id)

View File

@ -8,13 +8,18 @@ from typing import Any, Generic, TypeVar
from pylitterbot import FeederRobot, LitterRobot
from homeassistant.components.select import SelectEntity, SelectEntityDescription
from homeassistant.components.select import (
DOMAIN as PLATFORM,
SelectEntity,
SelectEntityDescription,
)
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import TIME_MINUTES
from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from .const import DOMAIN
from .entity import LitterRobotConfigEntity, _RobotT
from .entity import LitterRobotConfigEntity, _RobotT, async_update_unique_id
from .hub import LitterRobotHub
_CastTypeT = TypeVar("_CastTypeT", int, float)
@ -40,9 +45,10 @@ class RobotSelectEntityDescription(
LITTER_ROBOT_SELECT = RobotSelectEntityDescription[LitterRobot, int](
key="clean_cycle_wait_time_minutes",
key="cycle_delay",
name="Clean Cycle Wait Time Minutes",
icon="mdi:timer-outline",
unit_of_measurement=TIME_MINUTES,
current_fn=lambda robot: robot.clean_cycle_wait_time_minutes,
options_fn=lambda robot: robot.VALID_WAIT_TIMES,
select_fn=lambda robot, option: (robot.set_wait_time, int(option)),
@ -65,7 +71,7 @@ async def async_setup_entry(
) -> None:
"""Set up Litter-Robot selects using config entry."""
hub: LitterRobotHub = hass.data[DOMAIN][config_entry.entry_id]
async_add_entities(
entities: list[LitterRobotSelect] = list(
itertools.chain(
(
LitterRobotSelect(robot=robot, hub=hub, description=LITTER_ROBOT_SELECT)
@ -77,6 +83,8 @@ async def async_setup_entry(
),
)
)
async_update_unique_id(hass, PLATFORM, entities)
async_add_entities(entities)
class LitterRobotSelect(
@ -93,9 +101,7 @@ class LitterRobotSelect(
description: RobotSelectEntityDescription[_RobotT, _CastTypeT],
) -> None:
"""Initialize a Litter-Robot select entity."""
assert description.name
super().__init__(robot, description.name, hub)
self.entity_description = description
super().__init__(robot, hub, description)
options = self.entity_description.options_fn(self.robot)
self._attr_options = list(map(str, options))

View File

@ -9,6 +9,7 @@ from typing import Any, Generic, Union, cast
from pylitterbot import FeederRobot, LitterRobot, LitterRobot4, Robot
from homeassistant.components.sensor import (
DOMAIN as PLATFORM,
SensorDeviceClass,
SensorEntity,
SensorEntityDescription,
@ -20,7 +21,7 @@ from homeassistant.helpers.entity import EntityCategory
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from .const import DOMAIN
from .entity import LitterRobotEntity, _RobotT
from .entity import LitterRobotEntity, _RobotT, async_update_unique_id
from .hub import LitterRobotHub
@ -48,17 +49,6 @@ class LitterRobotSensorEntity(LitterRobotEntity[_RobotT], SensorEntity):
entity_description: RobotSensorEntityDescription[_RobotT]
def __init__(
self,
robot: _RobotT,
hub: LitterRobotHub,
description: RobotSensorEntityDescription[_RobotT],
) -> None:
"""Initialize a Litter-Robot sensor entity."""
assert description.name
super().__init__(robot, description.name, hub)
self.entity_description = description
@property
def native_value(self) -> float | datetime | str | None:
"""Return the state."""
@ -79,32 +69,32 @@ class LitterRobotSensorEntity(LitterRobotEntity[_RobotT], SensorEntity):
ROBOT_SENSOR_MAP: dict[type[Robot], list[RobotSensorEntityDescription]] = {
LitterRobot: [
RobotSensorEntityDescription[LitterRobot](
name="Waste Drawer",
key="waste_drawer_level",
name="Waste Drawer",
native_unit_of_measurement=PERCENTAGE,
icon_fn=lambda state: icon_for_gauge_level(state, 10),
),
RobotSensorEntityDescription[LitterRobot](
name="Sleep Mode Start Time",
key="sleep_mode_start_time",
name="Sleep Mode Start Time",
device_class=SensorDeviceClass.TIMESTAMP,
should_report=lambda robot: robot.sleep_mode_enabled,
),
RobotSensorEntityDescription[LitterRobot](
name="Sleep Mode End Time",
key="sleep_mode_end_time",
name="Sleep Mode End Time",
device_class=SensorDeviceClass.TIMESTAMP,
should_report=lambda robot: robot.sleep_mode_enabled,
),
RobotSensorEntityDescription[LitterRobot](
name="Last Seen",
key="last_seen",
name="Last Seen",
device_class=SensorDeviceClass.TIMESTAMP,
entity_category=EntityCategory.DIAGNOSTIC,
),
RobotSensorEntityDescription[LitterRobot](
name="Status Code",
key="status_code",
name="Status Code",
device_class="litterrobot__status_code",
entity_category=EntityCategory.DIAGNOSTIC,
),
@ -119,8 +109,8 @@ ROBOT_SENSOR_MAP: dict[type[Robot], list[RobotSensorEntityDescription]] = {
],
FeederRobot: [
RobotSensorEntityDescription[FeederRobot](
name="Food level",
key="food_level",
name="Food level",
native_unit_of_measurement=PERCENTAGE,
icon_fn=lambda state: icon_for_gauge_level(state, 10),
)
@ -135,10 +125,12 @@ 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 robot_type, entity_descriptions in ROBOT_SENSOR_MAP.items()
if isinstance(robot, robot_type)
for description in entity_descriptions
)
]
async_update_unique_id(hass, PLATFORM, entities)
async_add_entities(entities)

View File

@ -7,13 +7,17 @@ from typing import Any, Generic, Union
from pylitterbot import FeederRobot, LitterRobot
from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
from homeassistant.components.switch import (
DOMAIN as PLATFORM,
SwitchEntity,
SwitchEntityDescription,
)
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from .const import DOMAIN
from .entity import LitterRobotConfigEntity, _RobotT
from .entity import LitterRobotConfigEntity, _RobotT, async_update_unique_id
from .hub import LitterRobotHub
@ -51,17 +55,6 @@ class RobotSwitchEntity(LitterRobotConfigEntity[_RobotT], SwitchEntity):
entity_description: RobotSwitchEntityDescription[_RobotT]
def __init__(
self,
robot: _RobotT,
hub: LitterRobotHub,
description: RobotSwitchEntityDescription[_RobotT],
) -> None:
"""Initialize a Litter-Robot switch entity."""
assert description.name
super().__init__(robot, description.name, hub)
self.entity_description = description
@property
def is_on(self) -> bool | None:
"""Return true if switch is on."""
@ -93,9 +86,11 @@ async def async_setup_entry(
) -> None:
"""Set up Litter-Robot switches using config entry."""
hub: LitterRobotHub = hass.data[DOMAIN][entry.entry_id]
async_add_entities(
entities = [
RobotSwitchEntity(robot=robot, hub=hub, description=description)
for description in ROBOT_SWITCHES
for robot in hub.account.robots
if isinstance(robot, (LitterRobot, FeederRobot))
)
]
async_update_unique_id(hass, PLATFORM, entities)
async_add_entities(entities)

View File

@ -1,7 +1,6 @@
"""Support for Litter-Robot "Vacuum"."""
from __future__ import annotations
import logging
from typing import Any
from pylitterbot import LitterRobot
@ -9,11 +8,13 @@ from pylitterbot.enums import LitterBoxStatus
import voluptuous as vol
from homeassistant.components.vacuum import (
DOMAIN as PLATFORM,
STATE_CLEANING,
STATE_DOCKED,
STATE_ERROR,
STATE_PAUSED,
StateVacuumEntity,
StateVacuumEntityDescription,
VacuumEntityFeature,
)
from homeassistant.config_entries import ConfigEntry
@ -23,13 +24,9 @@ from homeassistant.helpers import config_validation as cv, entity_platform
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from .const import DOMAIN
from .entity import LitterRobotControlEntity
from .entity import LitterRobotControlEntity, async_update_unique_id
from .hub import LitterRobotHub
_LOGGER = logging.getLogger(__name__)
TYPE_LITTER_BOX = "Litter Box"
SERVICE_SET_SLEEP_MODE = "set_sleep_mode"
LITTER_BOX_STATUS_STATE_MAP = {
@ -44,6 +41,8 @@ LITTER_BOX_STATUS_STATE_MAP = {
LitterBoxStatus.OFF: STATE_OFF,
}
LITTER_BOX_ENTITY = StateVacuumEntityDescription("litter_box", name="Litter Box")
async def async_setup_entry(
hass: HomeAssistant,
@ -53,10 +52,12 @@ async def async_setup_entry(
"""Set up Litter-Robot cleaner using config entry."""
hub: LitterRobotHub = hass.data[DOMAIN][entry.entry_id]
async_add_entities(
LitterRobotCleaner(robot=robot, entity_type=TYPE_LITTER_BOX, hub=hub)
entities = [
LitterRobotCleaner(robot=robot, hub=hub, description=LITTER_BOX_ENTITY)
for robot in hub.litter_robots()
)
]
async_update_unique_id(hass, PLATFORM, entities)
async_add_entities(entities)
platform = entity_platform.async_get_current_platform()
platform.async_register_entity_service(

View File

@ -21,6 +21,7 @@ from homeassistant.components.vacuum import (
)
from homeassistant.const import ATTR_ENTITY_ID
from homeassistant.core import HomeAssistant
import homeassistant.helpers.entity_registry as er
from homeassistant.util.dt import utcnow
from .common import VACUUM_ENTITY_ID
@ -28,6 +29,9 @@ from .conftest import setup_integration
from tests.common import async_fire_time_changed
VACUUM_UNIQUE_ID_OLD = "LR3C012345-Litter Box"
VACUUM_UNIQUE_ID_NEW = "LR3C012345-litter_box"
COMPONENT_SERVICE_DOMAIN = {
SERVICE_SET_SLEEP_MODE: DOMAIN,
}
@ -35,6 +39,18 @@ COMPONENT_SERVICE_DOMAIN = {
async def test_vacuum(hass: HomeAssistant, mock_account: MagicMock) -> None:
"""Tests the vacuum entity was set up."""
ent_reg = er.async_get(hass)
# Create entity entry to migrate to new unique ID
ent_reg.async_get_or_create(
PLATFORM_DOMAIN,
DOMAIN,
VACUUM_UNIQUE_ID_OLD,
suggested_object_id=VACUUM_ENTITY_ID.replace(PLATFORM_DOMAIN, ""),
)
ent_reg_entry = ent_reg.async_get(VACUUM_ENTITY_ID)
assert ent_reg_entry.unique_id == VACUUM_UNIQUE_ID_OLD
await setup_integration(hass, mock_account, PLATFORM_DOMAIN)
assert hass.services.has_service(DOMAIN, SERVICE_SET_SLEEP_MODE)
@ -43,6 +59,9 @@ async def test_vacuum(hass: HomeAssistant, mock_account: MagicMock) -> None:
assert vacuum.state == STATE_DOCKED
assert vacuum.attributes["is_sleeping"] is False
ent_reg_entry = ent_reg.async_get(VACUUM_ENTITY_ID)
assert ent_reg_entry.unique_id == VACUUM_UNIQUE_ID_NEW
async def test_vacuum_status_when_sleeping(
hass: HomeAssistant, mock_account_with_sleeping_robot: MagicMock