From 64cef6e082dfac3e1fef050ebb237c53f1d97698 Mon Sep 17 00:00:00 2001 From: Robert Hillis Date: Fri, 21 Jun 2024 05:28:44 -0400 Subject: [PATCH] Store runtime data inside the config entry in Litter Robot (#119547) --- .../components/litterrobot/__init__.py | 29 +++++++++---------- .../components/litterrobot/binary_sensor.py | 8 ++--- .../components/litterrobot/button.py | 8 ++--- .../components/litterrobot/select.py | 7 ++--- .../components/litterrobot/sensor.py | 8 ++--- .../components/litterrobot/switch.py | 8 ++--- homeassistant/components/litterrobot/time.py | 8 ++--- .../components/litterrobot/update.py | 9 +++--- .../components/litterrobot/vacuum.py | 8 ++--- tests/components/litterrobot/test_init.py | 2 -- 10 files changed, 38 insertions(+), 57 deletions(-) diff --git a/homeassistant/components/litterrobot/__init__.py b/homeassistant/components/litterrobot/__init__.py index ec9849bbb89..3c55c4c4035 100644 --- a/homeassistant/components/litterrobot/__init__.py +++ b/homeassistant/components/litterrobot/__init__.py @@ -12,6 +12,8 @@ from homeassistant.helpers.device_registry import DeviceEntry from .const import DOMAIN from .hub import LitterRobotHub +type LitterRobotConfigEntry = ConfigEntry[LitterRobotHub] + PLATFORMS_BY_TYPE = { Robot: ( Platform.BINARY_SENSOR, @@ -37,40 +39,35 @@ def get_platforms_for_robots(robots: list[Robot]) -> set[Platform]: } -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry(hass: HomeAssistant, entry: LitterRobotConfigEntry) -> bool: """Set up Litter-Robot from a config entry.""" - hass.data.setdefault(DOMAIN, {}) - hub = hass.data[DOMAIN][entry.entry_id] = LitterRobotHub(hass, entry.data) + hub = LitterRobotHub(hass, entry.data) await hub.login(load_robots=True, subscribe_for_updates=True) + entry.runtime_data = hub if platforms := get_platforms_for_robots(hub.account.robots): await hass.config_entries.async_forward_entry_setups(entry, platforms) return True -async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_unload_entry( + hass: HomeAssistant, entry: LitterRobotConfigEntry +) -> bool: """Unload a config entry.""" - hub: LitterRobotHub = hass.data[DOMAIN][entry.entry_id] - await hub.account.disconnect() + await entry.runtime_data.account.disconnect() - platforms = get_platforms_for_robots(hub.account.robots) - unload_ok = await hass.config_entries.async_unload_platforms(entry, platforms) - - if unload_ok: - hass.data[DOMAIN].pop(entry.entry_id) - - return unload_ok + platforms = get_platforms_for_robots(entry.runtime_data.account.robots) + return await hass.config_entries.async_unload_platforms(entry, platforms) async def async_remove_config_entry_device( - hass: HomeAssistant, config_entry: ConfigEntry, device_entry: DeviceEntry + hass: HomeAssistant, entry: LitterRobotConfigEntry, device_entry: DeviceEntry ) -> bool: """Remove a config entry from a device.""" - hub: LitterRobotHub = hass.data[DOMAIN][config_entry.entry_id] return not any( identifier for identifier in device_entry.identifiers if identifier[0] == DOMAIN - for robot in hub.account.robots + for robot in entry.runtime_data.account.robots if robot.serial == identifier[1] ) diff --git a/homeassistant/components/litterrobot/binary_sensor.py b/homeassistant/components/litterrobot/binary_sensor.py index 2f44f44ed53..91113d6c094 100644 --- a/homeassistant/components/litterrobot/binary_sensor.py +++ b/homeassistant/components/litterrobot/binary_sensor.py @@ -13,14 +13,12 @@ from homeassistant.components.binary_sensor import ( BinarySensorEntity, BinarySensorEntityDescription, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from . import LitterRobotConfigEntry from .entity import LitterRobotEntity, _RobotT -from .hub import LitterRobotHub @dataclass(frozen=True) @@ -80,11 +78,11 @@ BINARY_SENSOR_MAP: dict[type[Robot], tuple[RobotBinarySensorEntityDescription, . async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: LitterRobotConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up Litter-Robot binary sensors using config entry.""" - hub: LitterRobotHub = hass.data[DOMAIN][entry.entry_id] + hub = entry.runtime_data async_add_entities( LitterRobotBinarySensorEntity(robot=robot, hub=hub, description=description) for robot in hub.account.robots diff --git a/homeassistant/components/litterrobot/button.py b/homeassistant/components/litterrobot/button.py index 02477e7fa03..6e6cc563c8e 100644 --- a/homeassistant/components/litterrobot/button.py +++ b/homeassistant/components/litterrobot/button.py @@ -10,23 +10,21 @@ from typing import Any, Generic from pylitterbot import FeederRobot, LitterRobot3 from homeassistant.components.button import ButtonEntity, ButtonEntityDescription -from homeassistant.config_entries import ConfigEntry from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from . import LitterRobotConfigEntry from .entity import LitterRobotEntity, _RobotT -from .hub import LitterRobotHub async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: LitterRobotConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up Litter-Robot cleaner using config entry.""" - hub: LitterRobotHub = hass.data[DOMAIN][entry.entry_id] + hub = entry.runtime_data entities: list[LitterRobotButtonEntity] = list( itertools.chain( ( diff --git a/homeassistant/components/litterrobot/select.py b/homeassistant/components/litterrobot/select.py index e7ecbada10d..948fad45a76 100644 --- a/homeassistant/components/litterrobot/select.py +++ b/homeassistant/components/litterrobot/select.py @@ -10,12 +10,11 @@ from pylitterbot import FeederRobot, LitterRobot, LitterRobot4, Robot from pylitterbot.robot.litterrobot4 import BrightnessLevel from homeassistant.components.select import SelectEntity, SelectEntityDescription -from homeassistant.config_entries import ConfigEntry from homeassistant.const import EntityCategory, UnitOfTime from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from . import LitterRobotConfigEntry from .entity import LitterRobotEntity, _RobotT from .hub import LitterRobotHub @@ -82,11 +81,11 @@ ROBOT_SELECT_MAP: dict[type[Robot], RobotSelectEntityDescription] = { async def async_setup_entry( hass: HomeAssistant, - config_entry: ConfigEntry, + entry: LitterRobotConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up Litter-Robot selects using config entry.""" - hub: LitterRobotHub = hass.data[DOMAIN][config_entry.entry_id] + hub = entry.runtime_data entities = [ LitterRobotSelectEntity(robot=robot, hub=hub, description=description) for robot in hub.account.robots diff --git a/homeassistant/components/litterrobot/sensor.py b/homeassistant/components/litterrobot/sensor.py index 1b4b7f78fdc..c110b89c7da 100644 --- a/homeassistant/components/litterrobot/sensor.py +++ b/homeassistant/components/litterrobot/sensor.py @@ -15,14 +15,12 @@ from homeassistant.components.sensor import ( SensorEntityDescription, SensorStateClass, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import PERCENTAGE, EntityCategory, UnitOfMass from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from . import LitterRobotConfigEntry from .entity import LitterRobotEntity, _RobotT -from .hub import LitterRobotHub def icon_for_gauge_level(gauge_level: int | None = None, offset: int = 0) -> str: @@ -157,11 +155,11 @@ ROBOT_SENSOR_MAP: dict[type[Robot], list[RobotSensorEntityDescription]] = { async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: LitterRobotConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up Litter-Robot sensors using config entry.""" - hub: LitterRobotHub = hass.data[DOMAIN][entry.entry_id] + hub = entry.runtime_data entities = [ LitterRobotSensorEntity(robot=robot, hub=hub, description=description) for robot in hub.account.robots diff --git a/homeassistant/components/litterrobot/switch.py b/homeassistant/components/litterrobot/switch.py index 60ca9b4d6c7..133fd897cc6 100644 --- a/homeassistant/components/litterrobot/switch.py +++ b/homeassistant/components/litterrobot/switch.py @@ -9,14 +9,12 @@ from typing import Any, Generic from pylitterbot import FeederRobot, LitterRobot from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription -from homeassistant.config_entries import ConfigEntry from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from . import LitterRobotConfigEntry from .entity import LitterRobotEntity, _RobotT -from .hub import LitterRobotHub @dataclass(frozen=True) @@ -68,11 +66,11 @@ class RobotSwitchEntity(LitterRobotEntity[_RobotT], SwitchEntity): async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: LitterRobotConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up Litter-Robot switches using config entry.""" - hub: LitterRobotHub = hass.data[DOMAIN][entry.entry_id] + hub = entry.runtime_data entities = [ RobotSwitchEntity(robot=robot, hub=hub, description=description) for description in ROBOT_SWITCHES diff --git a/homeassistant/components/litterrobot/time.py b/homeassistant/components/litterrobot/time.py index e2ada80b234..ace30d9f3a9 100644 --- a/homeassistant/components/litterrobot/time.py +++ b/homeassistant/components/litterrobot/time.py @@ -10,15 +10,13 @@ from typing import Any, Generic from pylitterbot import LitterRobot3 from homeassistant.components.time import TimeEntity, TimeEntityDescription -from homeassistant.config_entries import ConfigEntry from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback import homeassistant.util.dt as dt_util -from .const import DOMAIN +from . import LitterRobotConfigEntry from .entity import LitterRobotEntity, _RobotT -from .hub import LitterRobotHub @dataclass(frozen=True) @@ -52,11 +50,11 @@ LITTER_ROBOT_3_SLEEP_START = RobotTimeEntityDescription[LitterRobot3]( async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: LitterRobotConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up Litter-Robot cleaner using config entry.""" - hub: LitterRobotHub = hass.data[DOMAIN][entry.entry_id] + hub = entry.runtime_data async_add_entities( [ LitterRobotTimeEntity( diff --git a/homeassistant/components/litterrobot/update.py b/homeassistant/components/litterrobot/update.py index c4d1ada6080..1d3e1dff57c 100644 --- a/homeassistant/components/litterrobot/update.py +++ b/homeassistant/components/litterrobot/update.py @@ -13,13 +13,12 @@ from homeassistant.components.update import ( UpdateEntityDescription, UpdateEntityFeature, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN -from .entity import LitterRobotEntity, LitterRobotHub +from . import LitterRobotConfigEntry +from .entity import LitterRobotEntity SCAN_INTERVAL = timedelta(days=1) @@ -31,11 +30,11 @@ FIRMWARE_UPDATE_ENTITY = UpdateEntityDescription( async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: LitterRobotConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up Litter-Robot update platform.""" - hub: LitterRobotHub = hass.data[DOMAIN][entry.entry_id] + hub = entry.runtime_data entities = [ RobotUpdateEntity(robot=robot, hub=hub, description=FIRMWARE_UPDATE_ENTITY) for robot in hub.litter_robots() diff --git a/homeassistant/components/litterrobot/vacuum.py b/homeassistant/components/litterrobot/vacuum.py index d752609d7de..a1ed2ea600d 100644 --- a/homeassistant/components/litterrobot/vacuum.py +++ b/homeassistant/components/litterrobot/vacuum.py @@ -18,16 +18,14 @@ from homeassistant.components.vacuum import ( StateVacuumEntityDescription, VacuumEntityFeature, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import STATE_OFF from homeassistant.core import HomeAssistant from homeassistant.helpers import config_validation as cv, entity_platform from homeassistant.helpers.entity_platform import AddEntitiesCallback import homeassistant.util.dt as dt_util -from .const import DOMAIN +from . import LitterRobotConfigEntry from .entity import LitterRobotEntity -from .hub import LitterRobotHub SERVICE_SET_SLEEP_MODE = "set_sleep_mode" @@ -51,11 +49,11 @@ LITTER_BOX_ENTITY = StateVacuumEntityDescription( async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: LitterRobotConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up Litter-Robot cleaner using config entry.""" - hub: LitterRobotHub = hass.data[DOMAIN][entry.entry_id] + hub = entry.runtime_data entities = [ LitterRobotCleaner(robot=robot, hub=hub, description=LITTER_BOX_ENTITY) for robot in hub.litter_robots() diff --git a/tests/components/litterrobot/test_init.py b/tests/components/litterrobot/test_init.py index f4ad12aeb20..21b16097603 100644 --- a/tests/components/litterrobot/test_init.py +++ b/tests/components/litterrobot/test_init.py @@ -41,8 +41,6 @@ async def test_unload_entry(hass: HomeAssistant, mock_account: MagicMock) -> Non getattr(mock_account.robots[0], "start_cleaning").assert_called_once() assert await hass.config_entries.async_unload(entry.entry_id) - await hass.async_block_till_done() - assert hass.data[litterrobot.DOMAIN] == {} @pytest.mark.parametrize(