Use config entry runtime_data in aprilaire (#127079)

homeassistant/components/aprilaire/__init__.py

@@ -6,14 +6,12 @@ import logging
 
 from pyaprilaire.const import Attribute
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, CONF_PORT, EVENT_HOMEASSISTANT_STOP, Platform
 from homeassistant.core import Event, HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
 from homeassistant.helpers.device_registry import format_mac
 
-from .const import DOMAIN
-from .coordinator import AprilaireCoordinator
+from .coordinator import AprilaireConfigEntry, AprilaireCoordinator
 
 PLATFORMS: list[Platform] = [
     Platform.CLIMATE,
@@ -25,7 +23,7 @@ PLATFORMS: list[Platform] = [
 _LOGGER = logging.getLogger(__name__)
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: AprilaireConfigEntry) -> bool:
     """Set up a config entry for Aprilaire."""
 
     host = entry.data[CONF_HOST]
@@ -34,15 +32,16 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     coordinator = AprilaireCoordinator(hass, entry.unique_id, host, port)
     await coordinator.start_listen()
 
-    hass.data.setdefault(DOMAIN, {})[entry.unique_id] = coordinator
-
-    async def ready_callback(ready: bool):
+    async def ready_callback(ready: bool) -> None:
         if ready:
             mac_address = format_mac(coordinator.data[Attribute.MAC_ADDRESS])
 
             if mac_address != entry.unique_id:
                 raise ConfigEntryAuthFailed("Invalid MAC address")
 
+            entry.runtime_data = coordinator
+            entry.async_on_unload(coordinator.stop_listen)
+
             await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
             async def _async_close(_: Event) -> None:
@@ -63,12 +62,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: AprilaireConfigEntry) -> bool:
     """Unload a config entry."""
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-
-    if unload_ok:
-        coordinator: AprilaireCoordinator = hass.data[DOMAIN].pop(entry.unique_id)
-        coordinator.stop_listen()
-
-    return unload_ok
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)

homeassistant/components/aprilaire/climate.py

@@ -16,19 +16,17 @@ from homeassistant.components.climate import (
     HVACAction,
     HVACMode,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import PRECISION_HALVES, PRECISION_WHOLE, UnitOfTemperature
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .const import (
-    DOMAIN,
     FAN_CIRCULATE,
     PRESET_PERMANENT_HOLD,
     PRESET_TEMPORARY_HOLD,
     PRESET_VACATION,
 )
-from .coordinator import AprilaireCoordinator
+from .coordinator import AprilaireConfigEntry
 from .entity import BaseAprilaireEntity
 
 HVAC_MODE_MAP = {
@@ -64,14 +62,14 @@ FAN_MODE_MAP = {
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: AprilaireConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Add climates for passed config_entry in HA."""
 
-    coordinator: AprilaireCoordinator = hass.data[DOMAIN][config_entry.unique_id]
-
-    async_add_entities([AprilaireClimate(coordinator, config_entry.unique_id)])
+    async_add_entities(
+        [AprilaireClimate(config_entry.runtime_data, config_entry.unique_id)]
+    )
 
 
 class AprilaireClimate(BaseAprilaireEntity, ClimateEntity):

homeassistant/components/aprilaire/coordinator.py

@@ -9,6 +9,7 @@ from typing import Any
 import pyaprilaire.client
 from pyaprilaire.const import MODELS, Attribute, FunctionalDomain
 
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback
 import homeassistant.helpers.device_registry as dr
 from homeassistant.helpers.device_registry import DeviceInfo
@@ -22,6 +23,8 @@ WAIT_TIMEOUT = 30
 
 _LOGGER = logging.getLogger(__name__)
 
+type AprilaireConfigEntry = ConfigEntry[AprilaireCoordinator]
+
 
 class AprilaireCoordinator(BaseDataUpdateCoordinatorProtocol):
     """Coordinator for interacting with the thermostat."""
@@ -112,7 +115,7 @@ class AprilaireCoordinator(BaseDataUpdateCoordinatorProtocol):
         self.client.stop_listen()
 
     async def wait_for_ready(
-        self, ready_callback: Callable[[bool], Awaitable[bool]]
+        self, ready_callback: Callable[[bool], Awaitable[None]]
     ) -> bool:
         """Wait for the client to be ready."""
 

homeassistant/components/aprilaire/humidifier.py

@@ -14,13 +14,11 @@ from homeassistant.components.humidifier import (
     HumidifierEntity,
     HumidifierEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
 
-from .const import DOMAIN
-from .coordinator import AprilaireCoordinator
+from .coordinator import AprilaireConfigEntry, AprilaireCoordinator
 from .entity import BaseAprilaireEntity
 
 HUMIDIFIER_ACTION_MAP: dict[StateType, HumidifierAction] = {
@@ -41,12 +39,12 @@ DEHUMIDIFIER_ACTION_MAP: dict[StateType, HumidifierAction] = {
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: AprilaireConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Aprilaire humidifier devices."""
 
-    coordinator: AprilaireCoordinator = hass.data[DOMAIN][config_entry.unique_id]
+    coordinator = config_entry.runtime_data
 
     assert config_entry.unique_id is not None
 

homeassistant/components/aprilaire/select.py

@@ -9,12 +9,10 @@ from typing import cast
 from pyaprilaire.const import Attribute
 
 from homeassistant.components.select import SelectEntity, SelectEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
-from .coordinator import AprilaireCoordinator
+from .coordinator import AprilaireConfigEntry, AprilaireCoordinator
 from .entity import BaseAprilaireEntity
 
 AIR_CLEANING_EVENT_MAP = {0: "off", 3: "event_clean", 4: "allergies"}
@@ -25,12 +23,12 @@ FRESH_AIR_MODE_MAP = {0: "off", 1: "automatic"}
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: AprilaireConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Aprilaire select devices."""
 
-    coordinator: AprilaireCoordinator = hass.data[DOMAIN][config_entry.unique_id]
+    coordinator = config_entry.runtime_data
 
     assert config_entry.unique_id is not None
 

homeassistant/components/aprilaire/sensor.py

@@ -13,14 +13,12 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import PERCENTAGE, UnitOfTemperature
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
 
-from .const import DOMAIN
-from .coordinator import AprilaireCoordinator
+from .coordinator import AprilaireConfigEntry, AprilaireCoordinator
 from .entity import BaseAprilaireEntity
 
 DEHUMIDIFICATION_STATUS_MAP: dict[StateType, str] = {
@@ -76,12 +74,12 @@ def get_entities(
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: AprilaireConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Aprilaire sensor devices."""
 
-    coordinator: AprilaireCoordinator = hass.data[DOMAIN][config_entry.unique_id]
+    coordinator = config_entry.runtime_data
 
     assert config_entry.unique_id is not None