From dec03d4d25cd362eaac3603df2023e5cf50ccfc2 Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Mon, 30 Sep 2024 10:11:21 +0200
Subject: [PATCH] Use config entry runtime_data in awair (#127073)

---
 homeassistant/components/awair/__init__.py    | 28 ++++++++-----------
 homeassistant/components/awair/coordinator.py |  2 ++
 homeassistant/components/awair/sensor.py      |  9 +++---
 tests/components/awair/__init__.py            |  2 +-
 4 files changed, 18 insertions(+), 23 deletions(-)

diff --git a/homeassistant/components/awair/__init__.py b/homeassistant/components/awair/__init__.py
index aa810bf532b..528c658eff1 100644
--- a/homeassistant/components/awair/__init__.py
+++ b/homeassistant/components/awair/__init__.py
@@ -2,14 +2,13 @@
 
 from __future__ import annotations
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
-from .const import DOMAIN
 from .coordinator import (
     AwairCloudDataUpdateCoordinator,
+    AwairConfigEntry,
     AwairDataUpdateCoordinator,
     AwairLocalDataUpdateCoordinator,
 )
@@ -17,7 +16,9 @@ from .coordinator import (
 PLATFORMS = [Platform.SENSOR]
 
 
-async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
+async def async_setup_entry(
+    hass: HomeAssistant, config_entry: AwairConfigEntry
+) -> bool:
     """Set up Awair integration from a config entry."""
     session = async_get_clientsession(hass)
 
@@ -33,28 +34,21 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
 
     await coordinator.async_config_entry_first_refresh()
 
-    hass.data.setdefault(DOMAIN, {})
-    hass.data[DOMAIN][config_entry.entry_id] = coordinator
+    config_entry.runtime_data = coordinator
 
     await hass.config_entries.async_forward_entry_setups(config_entry, PLATFORMS)
 
     return True
 
 
-async def _async_update_listener(hass: HomeAssistant, entry: ConfigEntry) -> None:
+async def _async_update_listener(hass: HomeAssistant, entry: AwairConfigEntry) -> None:
     """Handle options update."""
-    coordinator: AwairLocalDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
-    if entry.title != coordinator.title:
+    if entry.title != entry.runtime_data.title:
         await hass.config_entries.async_reload(entry.entry_id)
 
 
-async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
+async def async_unload_entry(
+    hass: HomeAssistant, config_entry: AwairConfigEntry
+) -> bool:
     """Unload Awair configuration."""
-    unload_ok = await hass.config_entries.async_unload_platforms(
-        config_entry, PLATFORMS
-    )
-
-    if unload_ok:
-        hass.data[DOMAIN].pop(config_entry.entry_id)
-
-    return unload_ok
+    return await hass.config_entries.async_unload_platforms(config_entry, PLATFORMS)
diff --git a/homeassistant/components/awair/coordinator.py b/homeassistant/components/awair/coordinator.py
index b63efff7733..78f0d9d65f2 100644
--- a/homeassistant/components/awair/coordinator.py
+++ b/homeassistant/components/awair/coordinator.py
@@ -26,6 +26,8 @@ from .const import (
     UPDATE_INTERVAL_LOCAL,
 )
 
+type AwairConfigEntry = ConfigEntry[AwairDataUpdateCoordinator]
+
 
 @dataclass
 class AwairResult:
diff --git a/homeassistant/components/awair/sensor.py b/homeassistant/components/awair/sensor.py
index a62a15368be..c92009d9b1b 100644
--- a/homeassistant/components/awair/sensor.py
+++ b/homeassistant/components/awair/sensor.py
@@ -46,7 +46,7 @@ from .const import (
     ATTRIBUTION,
     DOMAIN,
 )
-from .coordinator import AwairDataUpdateCoordinator, AwairResult
+from .coordinator import AwairConfigEntry, AwairDataUpdateCoordinator
 
 DUST_ALIASES = [API_PM25, API_PM10]
 
@@ -132,15 +132,14 @@ SENSOR_TYPES_DUST: tuple[AwairSensorEntityDescription, ...] = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: AwairConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Awair sensor entity based on a config entry."""
-    coordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator = config_entry.runtime_data
     entities = []
 
-    data: list[AwairResult] = coordinator.data.values()
-    for result in data:
+    for result in coordinator.data.values():
         if result.air_data:
             entities.append(AwairSensor(result.device, coordinator, SENSOR_TYPE_SCORE))
             device_sensors = result.air_data.sensors.keys()
diff --git a/tests/components/awair/__init__.py b/tests/components/awair/__init__.py
index f93866263a2..0c0fd0eb522 100644
--- a/tests/components/awair/__init__.py
+++ b/tests/components/awair/__init__.py
@@ -2,7 +2,7 @@
 
 from unittest.mock import patch
 
-from homeassistant.components.awair import DOMAIN
+from homeassistant.components.awair.const import DOMAIN
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant