Use runtime_data in iqvia (#144984)

This commit is contained in:
epenet 2025-05-15 17:48:02 +02:00 committed by GitHub
parent d33a0f75fd
commit ace12958d1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 22 additions and 37 deletions

View File

@ -6,7 +6,6 @@ import asyncio
from pyiqvia import Client
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import Platform
from homeassistant.core import HomeAssistant
from homeassistant.exceptions import ConfigEntryNotReady
@ -14,7 +13,6 @@ from homeassistant.helpers import aiohttp_client
from .const import (
CONF_ZIP_CODE,
DOMAIN,
TYPE_ALLERGY_FORECAST,
TYPE_ALLERGY_INDEX,
TYPE_ALLERGY_OUTLOOK,
@ -23,14 +21,14 @@ from .const import (
TYPE_DISEASE_FORECAST,
TYPE_DISEASE_INDEX,
)
from .coordinator import IqviaUpdateCoordinator
from .coordinator import IqviaConfigEntry, IqviaUpdateCoordinator
DEFAULT_ATTRIBUTION = "Data provided by IQVIA™"
PLATFORMS = [Platform.SENSOR]
async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
async def async_setup_entry(hass: HomeAssistant, entry: IqviaConfigEntry) -> bool:
"""Set up IQVIA as config entry."""
if not entry.unique_id:
# If the config entry doesn't already have a unique ID, set one:
@ -75,18 +73,13 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
# Once we've successfully authenticated, we re-enable client request retries:
client.enable_request_retries()
hass.data.setdefault(DOMAIN, {})
hass.data[DOMAIN][entry.entry_id] = coordinators
entry.runtime_data = coordinators
await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
return True
async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
async def async_unload_entry(hass: HomeAssistant, entry: IqviaConfigEntry) -> bool:
"""Unload an OpenUV 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)

View File

@ -16,16 +16,18 @@ from .const import LOGGER
DEFAULT_SCAN_INTERVAL = timedelta(minutes=30)
type IqviaConfigEntry = ConfigEntry[dict[str, IqviaUpdateCoordinator]]
class IqviaUpdateCoordinator(DataUpdateCoordinator[dict[str, Any]]):
"""Custom DataUpdateCoordinator for IQVIA."""
config_entry: ConfigEntry
config_entry: IqviaConfigEntry
def __init__(
self,
hass: HomeAssistant,
config_entry: ConfigEntry,
config_entry: IqviaConfigEntry,
name: str,
update_method: Callable[[], Coroutine[Any, Any, dict[str, Any]]],
) -> None:

View File

@ -5,12 +5,11 @@ 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_UNIQUE_ID
from homeassistant.core import HomeAssistant
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
from .const import CONF_ZIP_CODE, DOMAIN
from .const import CONF_ZIP_CODE
from .coordinator import IqviaConfigEntry
CONF_CITY = "City"
CONF_DISPLAY_LOCATION = "DisplayLocation"
@ -33,19 +32,15 @@ TO_REDACT = {
async def async_get_config_entry_diagnostics(
hass: HomeAssistant, entry: ConfigEntry
hass: HomeAssistant, entry: IqviaConfigEntry
) -> dict[str, Any]:
"""Return diagnostics for a config entry."""
coordinators: dict[str, DataUpdateCoordinator[dict[str, Any]]] = hass.data[DOMAIN][
entry.entry_id
]
return {
"entry": async_redact_data(entry.as_dict(), TO_REDACT),
"data": async_redact_data(
{
data_type: coordinator.data
for data_type, coordinator in coordinators.items()
for data_type, coordinator in entry.runtime_data.items()
},
TO_REDACT,
),

View File

@ -2,28 +2,23 @@
from __future__ import annotations
from typing import Any
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import callback
from homeassistant.helpers.entity import EntityDescription
from homeassistant.helpers.update_coordinator import (
CoordinatorEntity,
DataUpdateCoordinator,
)
from homeassistant.helpers.update_coordinator import CoordinatorEntity
from .const import CONF_ZIP_CODE, DOMAIN, TYPE_ALLERGY_FORECAST, TYPE_ALLERGY_OUTLOOK
from .coordinator import IqviaConfigEntry, IqviaUpdateCoordinator
class IQVIAEntity(CoordinatorEntity[DataUpdateCoordinator[dict[str, Any]]]):
class IQVIAEntity(CoordinatorEntity[IqviaUpdateCoordinator]):
"""Define a base IQVIA entity."""
_attr_has_entity_name = True
def __init__(
self,
coordinator: DataUpdateCoordinator[dict[str, Any]],
entry: ConfigEntry,
coordinator: IqviaUpdateCoordinator,
entry: IqviaConfigEntry,
description: EntityDescription,
) -> None:
"""Initialize."""

View File

@ -12,7 +12,6 @@ from homeassistant.components.sensor import (
SensorEntityDescription,
SensorStateClass,
)
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import ATTR_STATE
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
@ -32,6 +31,7 @@ from .const import (
TYPE_DISEASE_INDEX,
TYPE_DISEASE_TODAY,
)
from .coordinator import IqviaConfigEntry
from .entity import IQVIAEntity
ATTR_ALLERGEN_AMOUNT = "allergen_amount"
@ -128,13 +128,13 @@ INDEX_SENSOR_DESCRIPTIONS = (
async def async_setup_entry(
hass: HomeAssistant,
entry: ConfigEntry,
entry: IqviaConfigEntry,
async_add_entities: AddConfigEntryEntitiesCallback,
) -> None:
"""Set up IQVIA sensors based on a config entry."""
sensors: list[ForecastSensor | IndexSensor] = [
ForecastSensor(
hass.data[DOMAIN][entry.entry_id][
entry.runtime_data[
API_CATEGORY_MAPPING.get(description.key, description.key)
],
entry,

View File

@ -4,7 +4,7 @@ from typing import Any
import pytest
from homeassistant.components.iqvia import CONF_ZIP_CODE, DOMAIN
from homeassistant.components.iqvia.const import CONF_ZIP_CODE, DOMAIN
from homeassistant.config_entries import SOURCE_USER
from homeassistant.core import HomeAssistant
from homeassistant.data_entry_flow import FlowResultType