Use runtime_data to store coordinator state (#140486)

Use runtime-data to save coordinator state
This commit is contained in:
Pieter Viljoen 2025-03-13 00:52:42 -07:00 committed by GitHub
parent f5412dd209
commit ffa6f42c0e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 44 additions and 48 deletions

View File

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

View File

@ -1,10 +1,13 @@
"""Constants for the PurpleAir integration.""" """Constants for the PurpleAir integration."""
import logging 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" DOMAIN: Final[str] = "purpleair"
CONF_SENSOR_INDICES = "sensor_indices"
CONF_SENSOR_INDICES: Final[str] = "sensor_indices"

View File

@ -46,12 +46,15 @@ SENSOR_FIELDS_TO_RETRIEVE = [
UPDATE_INTERVAL = timedelta(minutes=2) UPDATE_INTERVAL = timedelta(minutes=2)
type PurpleAirConfigEntry = ConfigEntry[PurpleAirDataUpdateCoordinator]
class PurpleAirDataUpdateCoordinator(DataUpdateCoordinator[GetSensorsResponse]): class PurpleAirDataUpdateCoordinator(DataUpdateCoordinator[GetSensorsResponse]):
"""Define a PurpleAir-specific coordinator.""" """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.""" """Initialize."""
self._api = API( self._api = API(
entry.data[CONF_API_KEY], entry.data[CONF_API_KEY],

View File

@ -5,7 +5,6 @@ from __future__ import annotations
from typing import Any from typing import Any
from homeassistant.components.diagnostics import async_redact_data from homeassistant.components.diagnostics import async_redact_data
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import ( from homeassistant.const import (
CONF_API_KEY, CONF_API_KEY,
CONF_LATITUDE, CONF_LATITUDE,
@ -14,8 +13,7 @@ from homeassistant.const import (
) )
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from .const import DOMAIN from .coordinator import PurpleAirConfigEntry
from .coordinator import PurpleAirDataUpdateCoordinator
CONF_TITLE = "title" CONF_TITLE = "title"
@ -30,14 +28,13 @@ TO_REDACT = {
async def async_get_config_entry_diagnostics( async def async_get_config_entry_diagnostics(
hass: HomeAssistant, entry: ConfigEntry hass: HomeAssistant, entry: PurpleAirConfigEntry
) -> dict[str, Any]: ) -> dict[str, Any]:
"""Return diagnostics for a config entry.""" """Return diagnostics for a config entry."""
coordinator: PurpleAirDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
return async_redact_data( return async_redact_data(
{ {
"entry": entry.as_dict(), "entry": entry.as_dict(),
"data": coordinator.data.model_dump(), "data": entry.runtime_data.data.model_dump(),
}, },
TO_REDACT, TO_REDACT,
) )

View File

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

View File

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

View File

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

View File

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