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

This commit is contained in:
Robert Hillis 2024-06-21 05:28:44 -04:00 committed by GitHub
parent 9c5879656c
commit 64cef6e082
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
10 changed files with 38 additions and 57 deletions

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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