From d0d1b303fd0a4dc5c5e1f9f16ae0327d9e3489ed Mon Sep 17 00:00:00 2001 From: Nathan Spencer Date: Thu, 1 Sep 2022 12:02:46 -0600 Subject: [PATCH] Code quality improvements for litterrobot integration (#77605) --- .../components/litterrobot/button.py | 52 +++++++++---------- .../components/litterrobot/entity.py | 50 ++++++++++++------ .../components/litterrobot/select.py | 20 ++++--- .../components/litterrobot/sensor.py | 32 +++++------- .../components/litterrobot/switch.py | 25 ++++----- .../components/litterrobot/vacuum.py | 19 +++---- tests/components/litterrobot/test_vacuum.py | 19 +++++++ 7 files changed, 123 insertions(+), 94 deletions(-) diff --git a/homeassistant/components/litterrobot/button.py b/homeassistant/components/litterrobot/button.py index 74c659fd474..81d9c65927e 100644 --- a/homeassistant/components/litterrobot/button.py +++ b/homeassistant/components/litterrobot/button.py @@ -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,21 +30,24 @@ 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( - ( - LitterRobotButtonEntity( - robot=robot, hub=hub, description=LITTER_ROBOT_BUTTON - ) - for robot in hub.litter_robots() - if isinstance(robot, LitterRobot3) - ), - ( - LitterRobotButtonEntity( - robot=robot, hub=hub, description=FEEDER_ROBOT_BUTTON - ) - for robot in hub.feeder_robots() - ), + entities: list[LitterRobotButtonEntity] = list( + itertools.chain( + ( + LitterRobotButtonEntity( + robot=robot, hub=hub, description=LITTER_ROBOT_BUTTON + ) + for robot in hub.litter_robots() + if isinstance(robot, LitterRobot3) + ), + ( + LitterRobotButtonEntity( + robot=robot, hub=hub, description=FEEDER_ROBOT_BUTTON + ) + 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) diff --git a/homeassistant/components/litterrobot/entity.py b/homeassistant/components/litterrobot/entity.py index 8471e007ce9..9716793f70e 100644 --- a/homeassistant/components/litterrobot/entity.py +++ b/homeassistant/components/litterrobot/entity.py @@ -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) diff --git a/homeassistant/components/litterrobot/select.py b/homeassistant/components/litterrobot/select.py index a18cd3b46b5..9ec784db8f2 100644 --- a/homeassistant/components/litterrobot/select.py +++ b/homeassistant/components/litterrobot/select.py @@ -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)) diff --git a/homeassistant/components/litterrobot/sensor.py b/homeassistant/components/litterrobot/sensor.py index 90bdfcbda73..c904335d23f 100644 --- a/homeassistant/components/litterrobot/sensor.py +++ b/homeassistant/components/litterrobot/sensor.py @@ -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) diff --git a/homeassistant/components/litterrobot/switch.py b/homeassistant/components/litterrobot/switch.py index 2f54ede38b8..779ee699b41 100644 --- a/homeassistant/components/litterrobot/switch.py +++ b/homeassistant/components/litterrobot/switch.py @@ -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) diff --git a/homeassistant/components/litterrobot/vacuum.py b/homeassistant/components/litterrobot/vacuum.py index 9a4b825045f..27cd3e6758a 100644 --- a/homeassistant/components/litterrobot/vacuum.py +++ b/homeassistant/components/litterrobot/vacuum.py @@ -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( diff --git a/tests/components/litterrobot/test_vacuum.py b/tests/components/litterrobot/test_vacuum.py index 02667bb8310..eb9a4c8c60b 100644 --- a/tests/components/litterrobot/test_vacuum.py +++ b/tests/components/litterrobot/test_vacuum.py @@ -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