diff --git a/homeassistant/components/husqvarna_automower/__init__.py b/homeassistant/components/husqvarna_automower/__init__.py index e4211e1078e..e62badd7e7c 100644 --- a/homeassistant/components/husqvarna_automower/__init__.py +++ b/homeassistant/components/husqvarna_automower/__init__.py @@ -12,7 +12,6 @@ from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady from homeassistant.helpers import aiohttp_client, config_entry_oauth2_flow from . import api -from .const import DOMAIN from .coordinator import AutomowerDataUpdateCoordinator _LOGGER = logging.getLogger(__name__) @@ -27,8 +26,10 @@ PLATFORMS: list[Platform] = [ Platform.SWITCH, ] +type AutomowerConfigEntry = ConfigEntry[AutomowerDataUpdateCoordinator] -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: + +async def async_setup_entry(hass: HomeAssistant, entry: AutomowerConfigEntry) -> bool: """Set up this integration using UI.""" implementation = ( await config_entry_oauth2_flow.async_get_config_entry_implementation( @@ -47,16 +48,17 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: if 400 <= err.status < 500: raise ConfigEntryAuthFailed from err raise ConfigEntryNotReady from err + coordinator = AutomowerDataUpdateCoordinator(hass, automower_api, entry) await coordinator.async_config_entry_first_refresh() + entry.runtime_data = coordinator + entry.async_create_background_task( hass, coordinator.client_listen(hass, entry, automower_api), "websocket_task", ) - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator - if "amc:api" not in entry.data["token"]["scope"]: # We raise ConfigEntryAuthFailed here because the websocket can't be used # without the scope. So only polling would be possible. @@ -66,9 +68,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: return True -async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_unload_entry(hass: HomeAssistant, entry: AutomowerConfigEntry) -> bool: """Handle unload of an entry.""" - unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS) - if unload_ok: - hass.data[DOMAIN].pop(entry.entry_id) - return unload_ok + return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) diff --git a/homeassistant/components/husqvarna_automower/binary_sensor.py b/homeassistant/components/husqvarna_automower/binary_sensor.py index e8e64e7ffc7..922f7deb99b 100644 --- a/homeassistant/components/husqvarna_automower/binary_sensor.py +++ b/homeassistant/components/husqvarna_automower/binary_sensor.py @@ -11,11 +11,10 @@ from homeassistant.components.binary_sensor import ( BinarySensorEntity, BinarySensorEntityDescription, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from . import AutomowerConfigEntry from .coordinator import AutomowerDataUpdateCoordinator from .entity import AutomowerBaseEntity @@ -49,10 +48,12 @@ BINARY_SENSOR_TYPES: tuple[AutomowerBinarySensorEntityDescription, ...] = ( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: AutomowerConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up binary sensor platform.""" - coordinator: AutomowerDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data async_add_entities( AutomowerBinarySensorEntity(mower_id, coordinator, description) for mower_id in coordinator.data diff --git a/homeassistant/components/husqvarna_automower/device_tracker.py b/homeassistant/components/husqvarna_automower/device_tracker.py index 780d1da76fb..74ad624a515 100644 --- a/homeassistant/components/husqvarna_automower/device_tracker.py +++ b/homeassistant/components/husqvarna_automower/device_tracker.py @@ -3,20 +3,21 @@ from typing import TYPE_CHECKING from homeassistant.components.device_tracker import SourceType, TrackerEntity -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from . import AutomowerConfigEntry from .coordinator import AutomowerDataUpdateCoordinator from .entity import AutomowerBaseEntity async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: AutomowerConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up device tracker platform.""" - coordinator: AutomowerDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data async_add_entities( AutomowerDeviceTrackerEntity(mower_id, coordinator) for mower_id in coordinator.data diff --git a/homeassistant/components/husqvarna_automower/diagnostics.py b/homeassistant/components/husqvarna_automower/diagnostics.py index f5677d4cb4b..658f6f94445 100644 --- a/homeassistant/components/husqvarna_automower/diagnostics.py +++ b/homeassistant/components/husqvarna_automower/diagnostics.py @@ -11,8 +11,8 @@ from homeassistant.const import CONF_ACCESS_TOKEN from homeassistant.core import HomeAssistant from homeassistant.helpers.device_registry import DeviceEntry +from . import AutomowerConfigEntry from .const import DOMAIN -from .coordinator import AutomowerDataUpdateCoordinator CONF_REFRESH_TOKEN = "refresh_token" POSITIONS = "positions" @@ -33,10 +33,10 @@ async def async_get_config_entry_diagnostics( async def async_get_device_diagnostics( - hass: HomeAssistant, entry: ConfigEntry, device: DeviceEntry + hass: HomeAssistant, entry: AutomowerConfigEntry, device: DeviceEntry ) -> dict[str, Any]: """Return diagnostics for a device entry.""" - coordinator: AutomowerDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data for identifier in device.identifiers: if identifier[0] == DOMAIN: if ( diff --git a/homeassistant/components/husqvarna_automower/lawn_mower.py b/homeassistant/components/husqvarna_automower/lawn_mower.py index 8ba9136364a..50333076308 100644 --- a/homeassistant/components/husqvarna_automower/lawn_mower.py +++ b/homeassistant/components/husqvarna_automower/lawn_mower.py @@ -10,12 +10,11 @@ from homeassistant.components.lawn_mower import ( LawnMowerEntity, LawnMowerEntityFeature, ) -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 AutomowerConfigEntry from .coordinator import AutomowerDataUpdateCoordinator from .entity import AutomowerControlEntity @@ -42,10 +41,12 @@ _LOGGER = logging.getLogger(__name__) async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: AutomowerConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up lawn mower platform.""" - coordinator: AutomowerDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data async_add_entities( AutomowerLawnMowerEntity(mower_id, coordinator) for mower_id in coordinator.data ) diff --git a/homeassistant/components/husqvarna_automower/number.py b/homeassistant/components/husqvarna_automower/number.py index 72c1d360da9..f6d55389195 100644 --- a/homeassistant/components/husqvarna_automower/number.py +++ b/homeassistant/components/husqvarna_automower/number.py @@ -11,14 +11,14 @@ from aioautomower.model import MowerAttributes, WorkArea from aioautomower.session import AutomowerSession from homeassistant.components.number import NumberEntity, NumberEntityDescription -from homeassistant.config_entries import ConfigEntry from homeassistant.const import PERCENTAGE, EntityCategory, Platform from homeassistant.core import HomeAssistant, callback from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers import entity_registry as er from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN, EXECUTION_TIME_DELAY +from . import AutomowerConfigEntry +from .const import EXECUTION_TIME_DELAY from .coordinator import AutomowerDataUpdateCoordinator from .entity import AutomowerControlEntity @@ -111,10 +111,12 @@ WORK_AREA_NUMBER_TYPES: tuple[AutomowerWorkAreaNumberEntityDescription, ...] = ( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: AutomowerConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up number platform.""" - coordinator: AutomowerDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data entities: list[NumberEntity] = [] for mower_id in coordinator.data: @@ -227,7 +229,7 @@ class AutomowerWorkAreaNumberEntity(AutomowerControlEntity, NumberEntity): def async_remove_entities( hass: HomeAssistant, coordinator: AutomowerDataUpdateCoordinator, - config_entry: ConfigEntry, + entry: AutomowerConfigEntry, mower_id: str, ) -> None: """Remove deleted work areas from Home Assistant.""" @@ -238,9 +240,7 @@ def async_remove_entities( for work_area_id in _work_areas: uid = f"{mower_id}_{work_area_id}_cutting_height_work_area" active_work_areas.add(uid) - for entity_entry in er.async_entries_for_config_entry( - entity_reg, config_entry.entry_id - ): + for entity_entry in er.async_entries_for_config_entry(entity_reg, entry.entry_id): if ( entity_entry.domain == Platform.NUMBER and (split := entity_entry.unique_id.split("_"))[0] == mower_id diff --git a/homeassistant/components/husqvarna_automower/select.py b/homeassistant/components/husqvarna_automower/select.py index 1baa90e2799..b647407581f 100644 --- a/homeassistant/components/husqvarna_automower/select.py +++ b/homeassistant/components/husqvarna_automower/select.py @@ -7,13 +7,12 @@ from aioautomower.exceptions import ApiException from aioautomower.model import HeadlightModes from homeassistant.components.select import SelectEntity -from homeassistant.config_entries import ConfigEntry from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from . import AutomowerConfigEntry from .coordinator import AutomowerDataUpdateCoordinator from .entity import AutomowerControlEntity @@ -29,10 +28,12 @@ HEADLIGHT_MODES: list = [ async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: AutomowerConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up select platform.""" - coordinator: AutomowerDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data async_add_entities( AutomowerSelectEntity(mower_id, coordinator) for mower_id in coordinator.data diff --git a/homeassistant/components/husqvarna_automower/sensor.py b/homeassistant/components/husqvarna_automower/sensor.py index 0ece16f8e83..4cc3bcf5e57 100644 --- a/homeassistant/components/husqvarna_automower/sensor.py +++ b/homeassistant/components/husqvarna_automower/sensor.py @@ -13,13 +13,12 @@ from homeassistant.components.sensor import ( SensorEntityDescription, SensorStateClass, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import PERCENTAGE, EntityCategory, UnitOfLength, UnitOfTime from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.typing import StateType -from .const import DOMAIN +from . import AutomowerConfigEntry from .coordinator import AutomowerDataUpdateCoordinator from .entity import AutomowerBaseEntity @@ -319,10 +318,12 @@ SENSOR_TYPES: tuple[AutomowerSensorEntityDescription, ...] = ( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: AutomowerConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up sensor platform.""" - coordinator: AutomowerDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data async_add_entities( AutomowerSensorEntity(mower_id, coordinator, description) for mower_id in coordinator.data diff --git a/homeassistant/components/husqvarna_automower/switch.py b/homeassistant/components/husqvarna_automower/switch.py index 4964c50eee5..fed2d3cfedc 100644 --- a/homeassistant/components/husqvarna_automower/switch.py +++ b/homeassistant/components/husqvarna_automower/switch.py @@ -14,14 +14,14 @@ from aioautomower.model import ( ) from homeassistant.components.switch import SwitchEntity -from homeassistant.config_entries import ConfigEntry from homeassistant.const import Platform from homeassistant.core import HomeAssistant, callback from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers import entity_registry as er from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN, EXECUTION_TIME_DELAY +from . import AutomowerConfigEntry +from .const import EXECUTION_TIME_DELAY from .coordinator import AutomowerDataUpdateCoordinator from .entity import AutomowerControlEntity @@ -44,10 +44,12 @@ ERROR_STATES = [ async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: AutomowerConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up switch platform.""" - coordinator: AutomowerDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data entities: list[SwitchEntity] = [] entities.extend( AutomowerScheduleSwitchEntity(mower_id, coordinator) @@ -196,7 +198,7 @@ class AutomowerStayOutZoneSwitchEntity(AutomowerControlEntity, SwitchEntity): def async_remove_entities( hass: HomeAssistant, coordinator: AutomowerDataUpdateCoordinator, - config_entry: ConfigEntry, + entry: AutomowerConfigEntry, mower_id: str, ) -> None: """Remove deleted stay-out-zones from Home Assistant.""" @@ -207,9 +209,7 @@ def async_remove_entities( for zones_uid in _zones.zones: uid = f"{mower_id}_{zones_uid}_stay_out_zones" active_zones.add(uid) - for entity_entry in er.async_entries_for_config_entry( - entity_reg, config_entry.entry_id - ): + for entity_entry in er.async_entries_for_config_entry(entity_reg, entry.entry_id): if ( entity_entry.domain == Platform.SWITCH and (split := entity_entry.unique_id.split("_"))[0] == mower_id