Store runtime data inside the config entry in Litter Robot (#119547)

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 = LitterRobotHub(hass, entry.data)
-    hub = hass.data[DOMAIN][entry.entry_id] = 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 entry.runtime_data.account.disconnect()
-    await hub.account.disconnect()
 
-    platforms = get_platforms_for_robots(hub.account.robots)
+    platforms = get_platforms_for_robots(entry.runtime_data.account.robots)
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, platforms)
+    return await hass.config_entries.async_unload_platforms(entry, platforms)
-
-    if unload_ok:
-        hass.data[DOMAIN].pop(entry.entry_id)
-
-    return unload_ok
 
 
 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]
     )

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

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

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

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

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

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(

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 . import LitterRobotConfigEntry
-from .entity import LitterRobotEntity, LitterRobotHub
+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()

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

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(