Migrate airvisual to use runtime_data (#125532)

* Migrate airvisual to use runtime_data

* Remove usedefault

* Adjust

homeassistant/components/airvisual/__init__.py

@@ -53,6 +53,8 @@ from .const import (
     LOGGER,
 )
 
+type AirVisualConfigEntry = ConfigEntry[DataUpdateCoordinator]
+
 # We use a raw string for the airvisual_pro domain (instead of importing the actual
 # constant) so that we can avoid listing it as a dependency:
 DOMAIN_AIRVISUAL_PRO = "airvisual_pro"
@@ -91,10 +93,9 @@ def async_get_cloud_coordinators_by_api_key(
 ) -> list[DataUpdateCoordinator]:
     """Get all DataUpdateCoordinator objects related to a particular API key."""
     return [
-        coordinator
+        entry.runtime_data
-        for entry_id, coordinator in hass.data[DOMAIN].items()
+        for entry in hass.config_entries.async_entries(DOMAIN)
-        if (entry := hass.config_entries.async_get_entry(entry_id))
+        if entry.data.get(CONF_API_KEY) == api_key and hasattr(entry, "runtime_data")
-        and entry.data.get(CONF_API_KEY) == api_key
     ]
 
 
@@ -172,7 +173,7 @@ def _standardize_geography_config_entry(
     hass.config_entries.async_update_entry(entry, **entry_updates)
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: AirVisualConfigEntry) -> bool:
     """Set up AirVisual as config entry."""
     if CONF_API_KEY not in entry.data:
         # If this is a migrated AirVisual Pro entry, there's no actual setup to do;
@@ -220,8 +221,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     entry.async_on_unload(entry.add_update_listener(async_reload_entry))
 
     await coordinator.async_config_entry_first_refresh()
-    hass.data.setdefault(DOMAIN, {})
+    entry.runtime_data = coordinator
-    hass.data[DOMAIN][entry.entry_id] = coordinator
 
     # Reassess the interval between 2 server requests
     async_sync_geo_coordinator_update_intervals(hass, entry.data[CONF_API_KEY])
@@ -231,7 +231,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return True
 
 
-async def async_migrate_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_migrate_entry(hass: HomeAssistant, entry: AirVisualConfigEntry) -> bool:
     """Migrate an old config entry."""
     version = entry.version
 
@@ -388,13 +388,11 @@ async def async_migrate_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: AirVisualConfigEntry) -> bool:
     """Unload an AirVisual config entry."""
     unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
 
-    if unload_ok:
+    if unload_ok and CONF_API_KEY in entry.data:
-        hass.data[DOMAIN].pop(entry.entry_id)
-        if CONF_API_KEY in entry.data:
         # Re-calculate the update interval period for any remaining consumers of
         # this API key:
         async_sync_geo_coordinator_update_intervals(hass, entry.data[CONF_API_KEY])
@@ -402,7 +400,7 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return unload_ok
 
 
-async def async_reload_entry(hass: HomeAssistant, entry: ConfigEntry) -> None:
+async def async_reload_entry(hass: HomeAssistant, entry: AirVisualConfigEntry) -> None:
     """Handle an options update."""
     await hass.config_entries.async_reload(entry.entry_id)
 

homeassistant/components/airvisual/diagnostics.py

@@ -5,7 +5,6 @@ from __future__ import annotations
 from typing import Any
 
 from homeassistant.components.diagnostics import async_redact_data
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     CONF_API_KEY,
     CONF_COUNTRY,
@@ -15,9 +14,9 @@ from homeassistant.const import (
     CONF_UNIQUE_ID,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from .const import CONF_CITY, DOMAIN
+from . import AirVisualConfigEntry
+from .const import CONF_CITY
 
 CONF_COORDINATES = "coordinates"
 CONF_TITLE = "title"
@@ -37,10 +36,10 @@ TO_REDACT = {
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, entry: ConfigEntry
+    hass: HomeAssistant, entry: AirVisualConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    coordinator: DataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator = entry.runtime_data
 
     return {
         "entry": async_redact_data(entry.as_dict(), TO_REDACT),

homeassistant/components/airvisual/sensor.py

@@ -26,8 +26,8 @@ from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from . import AirVisualEntity
+from . import AirVisualConfigEntry, AirVisualEntity
-from .const import CONF_CITY, DOMAIN
+from .const import CONF_CITY
 
 ATTR_CITY = "city"
 ATTR_COUNTRY = "country"
@@ -105,10 +105,12 @@ POLLUTANT_UNITS = {
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: AirVisualConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up AirVisual sensors based on a config entry."""
-    coordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator = entry.runtime_data
     async_add_entities(
         AirVisualGeographySensor(coordinator, entry, description, locale)
         for locale in GEOGRAPHY_SENSOR_LOCALES