Use runtime_data to store coordinator state (#140486)

Use runtime-data to save coordinator state

homeassistant/components/purpleair/__init__.py

@@ -2,37 +2,34 @@
 
 from __future__ import annotations
 
-from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 
-from .const import DOMAIN
-from .coordinator import PurpleAirDataUpdateCoordinator
-
-PLATFORMS = [Platform.SENSOR]
+from .const import PLATFORMS
+from .coordinator import PurpleAirConfigEntry, PurpleAirDataUpdateCoordinator
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
-    """Set up PurpleAir from a config entry."""
-    coordinator = PurpleAirDataUpdateCoordinator(hass, entry)
+async def async_setup_entry(hass: HomeAssistant, entry: PurpleAirConfigEntry) -> bool:
+    """Set up PurpleAir config entry."""
+    coordinator = PurpleAirDataUpdateCoordinator(
+        hass,
+        entry,
+    )
+    entry.runtime_data = coordinator
+
     await coordinator.async_config_entry_first_refresh()
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
-    entry.async_on_unload(entry.add_update_listener(async_handle_entry_update))
+    entry.async_on_unload(entry.add_update_listener(async_reload_entry))
 
     return True
 
 
-async def async_handle_entry_update(hass: HomeAssistant, entry: ConfigEntry) -> None:
-    """Handle an options update."""
+async def async_reload_entry(hass: HomeAssistant, entry: PurpleAirConfigEntry) -> None:
+    """Reload config entry."""
     await hass.config_entries.async_reload(entry.entry_id)
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
-    """Unload a config entry."""
-    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-        hass.data[DOMAIN].pop(entry.entry_id)
-
-    return unload_ok
+async def async_unload_entry(hass: HomeAssistant, entry: PurpleAirConfigEntry) -> bool:
+    """Unload config entry."""
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)

homeassistant/components/purpleair/const.py

@@ -1,10 +1,13 @@
 """Constants for the PurpleAir integration."""
 
 import logging
+from typing import Final
 
-DOMAIN = "purpleair"
+from homeassistant.const import Platform
 
-LOGGER = logging.getLogger(__package__)
+LOGGER: Final = logging.getLogger(__package__)
+PLATFORMS: Final = [Platform.SENSOR]
 
-CONF_READ_KEY = "read_key"
-CONF_SENSOR_INDICES = "sensor_indices"
+DOMAIN: Final[str] = "purpleair"
+
+CONF_SENSOR_INDICES: Final[str] = "sensor_indices"

homeassistant/components/purpleair/coordinator.py

@@ -46,12 +46,15 @@ SENSOR_FIELDS_TO_RETRIEVE = [
 UPDATE_INTERVAL = timedelta(minutes=2)
 
 
+type PurpleAirConfigEntry = ConfigEntry[PurpleAirDataUpdateCoordinator]
+
+
 class PurpleAirDataUpdateCoordinator(DataUpdateCoordinator[GetSensorsResponse]):
     """Define a PurpleAir-specific coordinator."""
 
-    config_entry: ConfigEntry
+    config_entry: PurpleAirConfigEntry
 
-    def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None:
+    def __init__(self, hass: HomeAssistant, entry: PurpleAirConfigEntry) -> None:
         """Initialize."""
         self._api = API(
             entry.data[CONF_API_KEY],

homeassistant/components/purpleair/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 .const import DOMAIN
-from .coordinator import PurpleAirDataUpdateCoordinator
+from .coordinator import PurpleAirConfigEntry
 
 CONF_TITLE = "title"
 
@@ -30,14 +28,13 @@ TO_REDACT = {
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, entry: ConfigEntry
+    hass: HomeAssistant, entry: PurpleAirConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    coordinator: PurpleAirDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
     return async_redact_data(
         {
             "entry": entry.as_dict(),
-            "data": coordinator.data.model_dump(),
+            "data": entry.runtime_data.data.model_dump(),
         },
         TO_REDACT,
     )

homeassistant/components/purpleair/entity.py

@@ -7,13 +7,12 @@ from typing import Any
 
 from aiopurpleair.models.sensors import SensorModel
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_LATITUDE, ATTR_LONGITUDE, CONF_SHOW_ON_MAP
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from .const import DOMAIN
-from .coordinator import PurpleAirDataUpdateCoordinator
+from .coordinator import PurpleAirConfigEntry, PurpleAirDataUpdateCoordinator
 
 
 class PurpleAirEntity(CoordinatorEntity[PurpleAirDataUpdateCoordinator]):
@@ -23,12 +22,11 @@ class PurpleAirEntity(CoordinatorEntity[PurpleAirDataUpdateCoordinator]):
 
     def __init__(
         self,
-        coordinator: PurpleAirDataUpdateCoordinator,
-        entry: ConfigEntry,
+        entry: PurpleAirConfigEntry,
         sensor_index: int,
     ) -> None:
         """Initialize."""
-        super().__init__(coordinator)
+        super().__init__(entry.runtime_data)
 
         self._sensor_index = sensor_index
 

homeassistant/components/purpleair/sensor.py

@@ -13,7 +13,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     CONCENTRATION_MICROGRAMS_PER_CUBIC_METER,
     PERCENTAGE,
@@ -27,8 +26,8 @@ from homeassistant.const import (
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from .const import CONF_SENSOR_INDICES, DOMAIN
-from .coordinator import PurpleAirDataUpdateCoordinator
+from .const import CONF_SENSOR_INDICES
+from .coordinator import PurpleAirConfigEntry
 from .entity import PurpleAirEntity
 
 CONCENTRATION_PARTICLES_PER_100_MILLILITERS = f"particles/100{UnitOfVolume.MILLILITERS}"
@@ -165,13 +164,12 @@ SENSOR_DESCRIPTIONS = [
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: PurpleAirConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up PurpleAir sensors based on a config entry."""
-    coordinator: PurpleAirDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
     async_add_entities(
-        PurpleAirSensorEntity(coordinator, entry, sensor_index, description)
+        PurpleAirSensorEntity(entry, sensor_index, description)
         for sensor_index in entry.options[CONF_SENSOR_INDICES]
         for description in SENSOR_DESCRIPTIONS
     )
@@ -184,13 +182,12 @@ class PurpleAirSensorEntity(PurpleAirEntity, SensorEntity):
 
     def __init__(
         self,
-        coordinator: PurpleAirDataUpdateCoordinator,
-        entry: ConfigEntry,
+        entry: PurpleAirConfigEntry,
         sensor_index: int,
         description: PurpleAirSensorEntityDescription,
     ) -> None:
         """Initialize."""
-        super().__init__(coordinator, entry, sensor_index)
+        super().__init__(entry, sensor_index)
 
         self._attr_unique_id = f"{self._sensor_index}-{description.key}"
         self.entity_description = description

tests/components/purpleair/conftest.py

@@ -8,7 +8,7 @@ from aiopurpleair.endpoints.sensors import NearbySensorResult
 from aiopurpleair.models.sensors import GetSensorsResponse
 import pytest
 
-from homeassistant.components.purpleair import DOMAIN
+from homeassistant.components.purpleair.const import DOMAIN
 from homeassistant.core import HomeAssistant
 
 from tests.common import MockConfigEntry, load_fixture
@@ -20,7 +20,7 @@ TEST_SENSOR_INDEX2 = 567890
 
 @pytest.fixture(name="api")
 def api_fixture(get_sensors_response: GetSensorsResponse) -> Mock:
-    """Define a fixture to return a mocked aiopurple API object."""
+    """Define a fixture to return a mocked aiopurpleair API object."""
     return Mock(
         async_check_api_key=AsyncMock(),
         get_map_url=Mock(return_value="http://example.com"),

tests/components/purpleair/test_config_flow.py

@@ -5,7 +5,7 @@ from unittest.mock import AsyncMock, patch
 from aiopurpleair.errors import InvalidApiKeyError, PurpleAirError
 import pytest
 
-from homeassistant.components.purpleair import DOMAIN
+from homeassistant.components.purpleair.const import DOMAIN
 from homeassistant.config_entries import SOURCE_USER
 from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResultType
@@ -288,6 +288,7 @@ async def test_options_remove_sensor(
     device_entry = device_registry.async_get_device(
         identifiers={(DOMAIN, str(TEST_SENSOR_INDEX1))}
     )
+    assert device_entry is not None
     result = await hass.config_entries.options.async_configure(
         result["flow_id"],
         user_input={"sensor_device_id": device_entry.id},