From 8e7026d6436beeee3b4a5c1811b42b5dea87d69d Mon Sep 17 00:00:00 2001
From: Marc Mueller <30130371+cdce8p@users.noreply.github.com>
Date: Thu, 2 May 2024 21:00:05 +0200
Subject: [PATCH] Use runtime_data for airnow (#116619)

---
 homeassistant/components/airnow/__init__.py    | 18 +++++++-----------
 homeassistant/components/airnow/diagnostics.py |  8 +++-----
 homeassistant/components/airnow/sensor.py      |  7 +++----
 3 files changed, 13 insertions(+), 20 deletions(-)

diff --git a/homeassistant/components/airnow/__init__.py b/homeassistant/components/airnow/__init__.py
index 8fba13164e7..5b06a25f13a 100644
--- a/homeassistant/components/airnow/__init__.py
+++ b/homeassistant/components/airnow/__init__.py
@@ -15,14 +15,16 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
-from .const import DOMAIN
+from .const import DOMAIN  # noqa: F401
 from .coordinator import AirNowDataUpdateCoordinator
 
 _LOGGER = logging.getLogger(__name__)
 PLATFORMS = [Platform.SENSOR]
 
+AirNowConfigEntry = ConfigEntry[AirNowDataUpdateCoordinator]
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+
+async def async_setup_entry(hass: HomeAssistant, entry: AirNowConfigEntry) -> bool:
     """Set up AirNow from a config entry."""
     api_key = entry.data[CONF_API_KEY]
     latitude = entry.data[CONF_LATITUDE]
@@ -44,8 +46,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     await coordinator.async_config_entry_first_refresh()
 
     # Store Entity and Initialize Platforms
-    hass.data.setdefault(DOMAIN, {})
-    hass.data[DOMAIN][entry.entry_id] = coordinator
+    entry.runtime_data = coordinator
 
     # Listen for option changes
     entry.async_on_unload(entry.add_update_listener(update_listener))
@@ -87,14 +88,9 @@ 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: AirNowConfigEntry) -> bool:
     """Unload a config 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)
 
 
 async def update_listener(hass: HomeAssistant, entry: ConfigEntry) -> None:
diff --git a/homeassistant/components/airnow/diagnostics.py b/homeassistant/components/airnow/diagnostics.py
index 39db915bef9..76cc35fb13c 100644
--- a/homeassistant/components/airnow/diagnostics.py
+++ b/homeassistant/components/airnow/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_LATITUDE,
@@ -14,8 +13,7 @@ from homeassistant.const import (
 )
 from homeassistant.core import HomeAssistant
 
-from . import AirNowDataUpdateCoordinator
-from .const import DOMAIN
+from . import AirNowConfigEntry
 
 ATTR_LATITUDE_CAP = "Latitude"
 ATTR_LONGITUDE_CAP = "Longitude"
@@ -40,10 +38,10 @@ TO_REDACT = {
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, entry: ConfigEntry
+    hass: HomeAssistant, entry: AirNowConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    coordinator: AirNowDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator = entry.runtime_data
 
     return async_redact_data(
         {
diff --git a/homeassistant/components/airnow/sensor.py b/homeassistant/components/airnow/sensor.py
index 1289b6c2b16..559478a69d3 100644
--- a/homeassistant/components/airnow/sensor.py
+++ b/homeassistant/components/airnow/sensor.py
@@ -13,7 +13,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     ATTR_TIME,
     CONCENTRATION_MICROGRAMS_PER_CUBIC_METER,
@@ -26,7 +25,7 @@ from homeassistant.helpers.typing import StateType
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 from homeassistant.util.dt import get_time_zone
 
-from . import AirNowDataUpdateCoordinator
+from . import AirNowConfigEntry, AirNowDataUpdateCoordinator
 from .const import (
     ATTR_API_AQI,
     ATTR_API_AQI_DESCRIPTION,
@@ -116,11 +115,11 @@ SENSOR_TYPES: tuple[AirNowEntityDescription, ...] = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: AirNowConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up AirNow sensor entities based on a config entry."""
-    coordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator = config_entry.runtime_data
 
     entities = [AirNowSensor(coordinator, description) for description in SENSOR_TYPES]