From 4766f48f47160fe29797a3b11acfab2cf51f6132 Mon Sep 17 00:00:00 2001 From: Joost Lekkerkerker Date: Wed, 12 Jun 2024 17:44:03 +0200 Subject: [PATCH] Migrate Airzone to runtime_data (#119494) --- homeassistant/components/airzone/__init__.py | 16 +++++++--------- .../components/airzone/binary_sensor.py | 8 +++++--- homeassistant/components/airzone/climate.py | 9 ++++++--- homeassistant/components/airzone/diagnostics.py | 8 +++----- homeassistant/components/airzone/entity.py | 3 ++- homeassistant/components/airzone/select.py | 8 +++++--- homeassistant/components/airzone/sensor.py | 9 ++++++--- homeassistant/components/airzone/water_heater.py | 9 ++++++--- tests/components/airzone/test_diagnostics.py | 1 - tests/components/airzone/util.py | 2 +- 10 files changed, 41 insertions(+), 32 deletions(-) diff --git a/homeassistant/components/airzone/__init__.py b/homeassistant/components/airzone/__init__.py index 1a65b92c3f4..754dfe90dce 100644 --- a/homeassistant/components/airzone/__init__.py +++ b/homeassistant/components/airzone/__init__.py @@ -17,7 +17,6 @@ from homeassistant.helpers import ( entity_registry as er, ) -from .const import DOMAIN from .coordinator import AirzoneUpdateCoordinator PLATFORMS: list[Platform] = [ @@ -30,10 +29,12 @@ PLATFORMS: list[Platform] = [ _LOGGER = logging.getLogger(__name__) +type AirzoneConfigEntry = ConfigEntry[AirzoneUpdateCoordinator] + async def _async_migrate_unique_ids( hass: HomeAssistant, - entry: ConfigEntry, + entry: AirzoneConfigEntry, coordinator: AirzoneUpdateCoordinator, ) -> None: """Migrate entities when the mac address gets discovered.""" @@ -71,7 +72,7 @@ async def _async_migrate_unique_ids( await er.async_migrate_entries(hass, entry.entry_id, _async_migrator) -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry(hass: HomeAssistant, entry: AirzoneConfigEntry) -> bool: """Set up Airzone from a config entry.""" options = ConnectionOptions( entry.data[CONF_HOST], @@ -84,16 +85,13 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: await coordinator.async_config_entry_first_refresh() await _async_migrate_unique_ids(hass, entry, coordinator) - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator + entry.runtime_data = coordinator 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: AirzoneConfigEntry) -> bool: """Unload a config entry.""" - if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS): - 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/airzone/binary_sensor.py b/homeassistant/components/airzone/binary_sensor.py index e25751f2a47..20878c08b82 100644 --- a/homeassistant/components/airzone/binary_sensor.py +++ b/homeassistant/components/airzone/binary_sensor.py @@ -25,7 +25,7 @@ from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from . import AirzoneConfigEntry from .coordinator import AirzoneUpdateCoordinator from .entity import AirzoneEntity, AirzoneSystemEntity, AirzoneZoneEntity @@ -75,10 +75,12 @@ ZONE_BINARY_SENSOR_TYPES: Final[tuple[AirzoneBinarySensorEntityDescription, ...] async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: AirzoneConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Add Airzone binary sensors from a config_entry.""" - coordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data binary_sensors: list[AirzoneBinarySensor] = [ AirzoneSystemBinarySensor( diff --git a/homeassistant/components/airzone/climate.py b/homeassistant/components/airzone/climate.py index f5b42c4ccbd..33c84b67501 100644 --- a/homeassistant/components/airzone/climate.py +++ b/homeassistant/components/airzone/climate.py @@ -50,7 +50,8 @@ from homeassistant.core import HomeAssistant, callback from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import API_TEMPERATURE_STEP, DOMAIN, TEMP_UNIT_LIB_TO_HASS +from . import AirzoneConfigEntry +from .const import API_TEMPERATURE_STEP, TEMP_UNIT_LIB_TO_HASS from .coordinator import AirzoneUpdateCoordinator from .entity import AirzoneZoneEntity @@ -97,10 +98,12 @@ HVAC_MODE_HASS_TO_LIB: Final[dict[HVACMode, OperationMode]] = { async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: AirzoneConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Add Airzone sensors from a config_entry.""" - coordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data async_add_entities( AirzoneClimate( coordinator, diff --git a/homeassistant/components/airzone/diagnostics.py b/homeassistant/components/airzone/diagnostics.py index 8c75302d692..6c75b750eaf 100644 --- a/homeassistant/components/airzone/diagnostics.py +++ b/homeassistant/components/airzone/diagnostics.py @@ -7,12 +7,10 @@ from typing import Any from aioairzone.const import API_MAC, AZD_MAC from homeassistant.components.diagnostics.util import async_redact_data -from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_UNIQUE_ID from homeassistant.core import HomeAssistant -from .const import DOMAIN -from .coordinator import AirzoneUpdateCoordinator +from . import AirzoneConfigEntry TO_REDACT_API = [ API_MAC, @@ -28,10 +26,10 @@ TO_REDACT_COORD = [ async def async_get_config_entry_diagnostics( - hass: HomeAssistant, config_entry: ConfigEntry + hass: HomeAssistant, config_entry: AirzoneConfigEntry ) -> dict[str, Any]: """Return diagnostics for a config entry.""" - coordinator: AirzoneUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id] + coordinator = config_entry.runtime_data return { "api_data": async_redact_data(coordinator.airzone.raw_data(), TO_REDACT_API), diff --git a/homeassistant/components/airzone/entity.py b/homeassistant/components/airzone/entity.py index b360db61897..61f79eabf52 100644 --- a/homeassistant/components/airzone/entity.py +++ b/homeassistant/components/airzone/entity.py @@ -31,6 +31,7 @@ from homeassistant.helpers import device_registry as dr from homeassistant.helpers.device_registry import DeviceInfo from homeassistant.helpers.update_coordinator import CoordinatorEntity +from . import AirzoneConfigEntry from .const import DOMAIN, MANUFACTURER from .coordinator import AirzoneUpdateCoordinator @@ -53,7 +54,7 @@ class AirzoneSystemEntity(AirzoneEntity): def __init__( self, coordinator: AirzoneUpdateCoordinator, - entry: ConfigEntry, + entry: AirzoneConfigEntry, system_data: dict[str, Any], ) -> None: """Initialize.""" diff --git a/homeassistant/components/airzone/select.py b/homeassistant/components/airzone/select.py index 6e92394bb05..8ffe86851b8 100644 --- a/homeassistant/components/airzone/select.py +++ b/homeassistant/components/airzone/select.py @@ -22,7 +22,7 @@ from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN +from . import AirzoneConfigEntry from .coordinator import AirzoneUpdateCoordinator from .entity import AirzoneEntity, AirzoneZoneEntity @@ -79,10 +79,12 @@ ZONE_SELECT_TYPES: Final[tuple[AirzoneSelectDescription, ...]] = ( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: AirzoneConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Add Airzone sensors from a config_entry.""" - coordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data async_add_entities( AirzoneZoneSelect( diff --git a/homeassistant/components/airzone/sensor.py b/homeassistant/components/airzone/sensor.py index e2f9eabc6f6..7cba0dc515c 100644 --- a/homeassistant/components/airzone/sensor.py +++ b/homeassistant/components/airzone/sensor.py @@ -30,7 +30,8 @@ from homeassistant.const import ( from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN, TEMP_UNIT_LIB_TO_HASS +from . import AirzoneConfigEntry +from .const import TEMP_UNIT_LIB_TO_HASS from .coordinator import AirzoneUpdateCoordinator from .entity import ( AirzoneEntity, @@ -77,10 +78,12 @@ ZONE_SENSOR_TYPES: Final[tuple[SensorEntityDescription, ...]] = ( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: AirzoneConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Add Airzone sensors from a config_entry.""" - coordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data sensors: list[AirzoneSensor] = [ AirzoneZoneSensor( diff --git a/homeassistant/components/airzone/water_heater.py b/homeassistant/components/airzone/water_heater.py index 4e502776185..ed1c2069c27 100644 --- a/homeassistant/components/airzone/water_heater.py +++ b/homeassistant/components/airzone/water_heater.py @@ -30,7 +30,8 @@ from homeassistant.const import ATTR_TEMPERATURE, STATE_OFF from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddEntitiesCallback -from .const import DOMAIN, TEMP_UNIT_LIB_TO_HASS +from . import AirzoneConfigEntry +from .const import TEMP_UNIT_LIB_TO_HASS from .coordinator import AirzoneUpdateCoordinator from .entity import AirzoneHotWaterEntity @@ -56,10 +57,12 @@ OPERATION_MODE_TO_DHW_PARAMS: Final[dict[str, dict[str, Any]]] = { async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: AirzoneConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Add Airzone sensors from a config_entry.""" - coordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data if AZD_HOT_WATER in coordinator.data: async_add_entities([AirzoneWaterHeater(coordinator, entry)]) diff --git a/tests/components/airzone/test_diagnostics.py b/tests/components/airzone/test_diagnostics.py index b64f346f27e..6a03b9f1985 100644 --- a/tests/components/airzone/test_diagnostics.py +++ b/tests/components/airzone/test_diagnostics.py @@ -26,7 +26,6 @@ async def test_config_entry_diagnostics( ) -> None: """Test config entry diagnostics.""" await async_init_integration(hass) - assert hass.data[DOMAIN] config_entry = hass.config_entries.async_entries(DOMAIN)[0] with patch( diff --git a/tests/components/airzone/util.py b/tests/components/airzone/util.py index c5c2d5972d4..6e3e0eccc8f 100644 --- a/tests/components/airzone/util.py +++ b/tests/components/airzone/util.py @@ -55,7 +55,7 @@ from aioairzone.const import ( API_ZONE_ID, ) -from homeassistant.components.airzone import DOMAIN +from homeassistant.components.airzone.const import DOMAIN from homeassistant.const import CONF_HOST, CONF_ID, CONF_PORT from homeassistant.core import HomeAssistant