Store Airly runtime data in config entry (#117031)

* Store runtime data in config entry

* Fix tests

---------

Co-authored-by: Maciej Bieniek <478555+bieniu@users.noreply.github.com>

homeassistant/components/airly/__init__.py

@@ -19,8 +19,10 @@ PLATFORMS = [Platform.SENSOR]
 
 _LOGGER = logging.getLogger(__name__)
 
+AirlyConfigEntry = ConfigEntry[AirlyDataUpdateCoordinator]
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+
+async def async_setup_entry(hass: HomeAssistant, entry: AirlyConfigEntry) -> bool:
     """Set up Airly as config entry."""
     api_key = entry.data[CONF_API_KEY]
     latitude = entry.data[CONF_LATITUDE]
@@ -62,8 +64,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     )
     await coordinator.async_config_entry_first_refresh()
 
-    hass.data.setdefault(DOMAIN, {})
+    entry.runtime_data = coordinator
-    hass.data[DOMAIN][entry.entry_id] = coordinator
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
@@ -79,11 +80,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: AirlyConfigEntry) -> bool:
     """Unload a config entry."""
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-
-    if unload_ok:
-        hass.data[DOMAIN].pop(entry.entry_id)
-
-    return unload_ok

homeassistant/components/airly/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,17 +13,16 @@ from homeassistant.const import (
 )
 from homeassistant.core import HomeAssistant
 
-from . import AirlyDataUpdateCoordinator
+from . import AirlyConfigEntry
-from .const import DOMAIN
 
 TO_REDACT = {CONF_API_KEY, CONF_LATITUDE, CONF_LONGITUDE, CONF_UNIQUE_ID}
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, config_entry: ConfigEntry
+    hass: HomeAssistant, config_entry: AirlyConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    coordinator: AirlyDataUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator = config_entry.runtime_data
 
     return {
         "config_entry": async_redact_data(config_entry.as_dict(), TO_REDACT),

homeassistant/components/airly/sensor.py

@@ -12,7 +12,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     CONCENTRATION_MICROGRAMS_PER_CUBIC_METER,
     CONF_NAME,
@@ -25,7 +24,7 @@ from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
-from . import AirlyDataUpdateCoordinator
+from . import AirlyConfigEntry, AirlyDataUpdateCoordinator
 from .const import (
     ATTR_ADVICE,
     ATTR_API_ADVICE,
@@ -174,12 +173,14 @@ SENSOR_TYPES: tuple[AirlySensorEntityDescription, ...] = (
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: AirlyConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Airly sensor entities based on a config entry."""
     name = entry.data[CONF_NAME]
 
-    coordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator = entry.runtime_data
 
     async_add_entities(
         (

homeassistant/components/airly/system_health.py

@@ -9,6 +9,7 @@ from airly import Airly
 from homeassistant.components import system_health
 from homeassistant.core import HomeAssistant, callback
 
+from . import AirlyConfigEntry
 from .const import DOMAIN
 
 
@@ -22,8 +23,10 @@ def async_register(
 
 async def system_health_info(hass: HomeAssistant) -> dict[str, Any]:
     """Get info for the info page."""
-    requests_remaining = list(hass.data[DOMAIN].values())[0].airly.requests_remaining
+    config_entry: AirlyConfigEntry = hass.config_entries.async_entries(DOMAIN)[0]
-    requests_per_day = list(hass.data[DOMAIN].values())[0].airly.requests_per_day
+
+    requests_remaining = config_entry.runtime_data.airly.requests_remaining
+    requests_per_day = config_entry.runtime_data.airly.requests_per_day
 
     return {
         "can_reach_server": system_health.async_check_can_reach_url(

tests/components/airly/__init__.py

@@ -8,6 +8,10 @@ API_NEAREST_URL = "https://airapi.airly.eu/v2/measurements/nearest?lat=123.00000
 API_POINT_URL = (
     "https://airapi.airly.eu/v2/measurements/point?lat=123.000000&lng=456.000000"
 )
+HEADERS = {
+    "X-RateLimit-Limit-day": "100",
+    "X-RateLimit-Remaining-day": "42",
+}
 
 
 async def init_integration(hass, aioclient_mock) -> MockConfigEntry:
@@ -25,7 +29,9 @@ async def init_integration(hass, aioclient_mock) -> MockConfigEntry:
         },
     )
 
-    aioclient_mock.get(API_POINT_URL, text=load_fixture("valid_station.json", "airly"))
+    aioclient_mock.get(
+        API_POINT_URL, text=load_fixture("valid_station.json", DOMAIN), headers=HEADERS
+    )
     entry.add_to_hass(hass)
     await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()

tests/components/airly/test_system_health.py

@@ -1,7 +1,6 @@
 """Test Airly system health."""
 
 import asyncio
-from unittest.mock import Mock
 
 from aiohttp import ClientError
 
@@ -9,6 +8,8 @@ from homeassistant.components.airly.const import DOMAIN
 from homeassistant.core import HomeAssistant
 from homeassistant.setup import async_setup_component
 
+from . import init_integration
+
 from tests.common import get_system_health_info
 from tests.test_util.aiohttp import AiohttpClientMocker
 
@@ -18,19 +19,11 @@ async def test_airly_system_health(
 ) -> None:
     """Test Airly system health."""
     aioclient_mock.get("https://airapi.airly.eu/v2/", text="")
-    hass.config.components.add(DOMAIN)
+
+    await init_integration(hass, aioclient_mock)
     assert await async_setup_component(hass, "system_health", {})
     await hass.async_block_till_done()
 
-    hass.data[DOMAIN] = {}
-    hass.data[DOMAIN]["0123xyz"] = Mock(
-        airly=Mock(
-            AIRLY_API_URL="https://airapi.airly.eu/v2/",
-            requests_remaining=42,
-            requests_per_day=100,
-        )
-    )
-
     info = await get_system_health_info(hass, DOMAIN)
 
     for key, val in info.items():
@@ -47,19 +40,11 @@ async def test_airly_system_health_fail(
 ) -> None:
     """Test Airly system health."""
     aioclient_mock.get("https://airapi.airly.eu/v2/", exc=ClientError)
-    hass.config.components.add(DOMAIN)
+
+    await init_integration(hass, aioclient_mock)
     assert await async_setup_component(hass, "system_health", {})
     await hass.async_block_till_done()
 
-    hass.data[DOMAIN] = {}
-    hass.data[DOMAIN]["0123xyz"] = Mock(
-        airly=Mock(
-            AIRLY_API_URL="https://airapi.airly.eu/v2/",
-            requests_remaining=0,
-            requests_per_day=1000,
-        )
-    )
-
     info = await get_system_health_info(hass, DOMAIN)
 
     for key, val in info.items():
@@ -67,5 +52,5 @@ async def test_airly_system_health_fail(
             info[key] = await val
 
     assert info["can_reach_server"] == {"type": "failed", "error": "unreachable"}
-    assert info["requests_remaining"] == 0
+    assert info["requests_remaining"] == 42
-    assert info["requests_per_day"] == 1000
+    assert info["requests_per_day"] == 100