diff --git a/homeassistant/components/iqvia/__init__.py b/homeassistant/components/iqvia/__init__.py index 049c23965b1..ad8b78bf9e3 100644 --- a/homeassistant/components/iqvia/__init__.py +++ b/homeassistant/components/iqvia/__init__.py @@ -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) diff --git a/homeassistant/components/iqvia/coordinator.py b/homeassistant/components/iqvia/coordinator.py index 420cbadbefa..ef926d1112d 100644 --- a/homeassistant/components/iqvia/coordinator.py +++ b/homeassistant/components/iqvia/coordinator.py @@ -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: diff --git a/homeassistant/components/iqvia/diagnostics.py b/homeassistant/components/iqvia/diagnostics.py index 64827f183ff..953d42eafc2 100644 --- a/homeassistant/components/iqvia/diagnostics.py +++ b/homeassistant/components/iqvia/diagnostics.py @@ -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, ), diff --git a/homeassistant/components/iqvia/entity.py b/homeassistant/components/iqvia/entity.py index e77c0f7e32a..1964a7cb039 100644 --- a/homeassistant/components/iqvia/entity.py +++ b/homeassistant/components/iqvia/entity.py @@ -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.""" diff --git a/homeassistant/components/iqvia/sensor.py b/homeassistant/components/iqvia/sensor.py index 64492c634e9..c0401b27368 100644 --- a/homeassistant/components/iqvia/sensor.py +++ b/homeassistant/components/iqvia/sensor.py @@ -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, diff --git a/tests/components/iqvia/test_config_flow.py b/tests/components/iqvia/test_config_flow.py index 22f473a3fb5..9a973ebe49c 100644 --- a/tests/components/iqvia/test_config_flow.py +++ b/tests/components/iqvia/test_config_flow.py @@ -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