From 67223b2a2dbe6f9163f9a35ad614c8dfed52f616 Mon Sep 17 00:00:00 2001 From: Robert Hillis Date: Tue, 18 Jun 2024 03:13:24 -0400 Subject: [PATCH] Store runtime data inside the config entry in Lidarr (#119548) --- homeassistant/components/lidarr/__init__.py | 40 ++++++++++++------- .../components/lidarr/config_flow.py | 5 ++- homeassistant/components/lidarr/sensor.py | 12 ++---- 3 files changed, 32 insertions(+), 25 deletions(-) diff --git a/homeassistant/components/lidarr/__init__.py b/homeassistant/components/lidarr/__init__.py index acfb8f30f30..ee2369a6bc4 100644 --- a/homeassistant/components/lidarr/__init__.py +++ b/homeassistant/components/lidarr/__init__.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import Any +from dataclasses import dataclass, fields from aiopyarr.lidarr_client import LidarrClient from aiopyarr.models.host_configuration import PyArrHostConfiguration @@ -25,10 +25,22 @@ from .coordinator import ( WantedDataUpdateCoordinator, ) +type LidarrConfigEntry = ConfigEntry[LidarrData] + PLATFORMS = [Platform.SENSOR] -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +@dataclass(kw_only=True, slots=True) +class LidarrData: + """Lidarr data type.""" + + disk_space: DiskSpaceDataUpdateCoordinator + queue: QueueDataUpdateCoordinator + status: StatusDataUpdateCoordinator + wanted: WantedDataUpdateCoordinator + + +async def async_setup_entry(hass: HomeAssistant, entry: LidarrConfigEntry) -> bool: """Set up Lidarr from a config entry.""" host_configuration = PyArrHostConfiguration( api_token=entry.data[CONF_API_KEY], @@ -40,31 +52,29 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: session=async_get_clientsession(hass, host_configuration.verify_ssl), request_timeout=60, ) - coordinators: dict[str, LidarrDataUpdateCoordinator[Any]] = { - "disk_space": DiskSpaceDataUpdateCoordinator(hass, host_configuration, lidarr), - "queue": QueueDataUpdateCoordinator(hass, host_configuration, lidarr), - "status": StatusDataUpdateCoordinator(hass, host_configuration, lidarr), - "wanted": WantedDataUpdateCoordinator(hass, host_configuration, lidarr), - } + data = LidarrData( + disk_space=DiskSpaceDataUpdateCoordinator(hass, host_configuration, lidarr), + queue=QueueDataUpdateCoordinator(hass, host_configuration, lidarr), + status=StatusDataUpdateCoordinator(hass, host_configuration, lidarr), + wanted=WantedDataUpdateCoordinator(hass, host_configuration, lidarr), + ) # Temporary, until we add diagnostic entities _version = None - for coordinator in coordinators.values(): + for field in fields(data): + coordinator = getattr(data, field.name) await coordinator.async_config_entry_first_refresh() if isinstance(coordinator, StatusDataUpdateCoordinator): _version = coordinator.data coordinator.system_version = _version - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinators + entry.runtime_data = data 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: LidarrConfigEntry) -> 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 + return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) class LidarrEntity(CoordinatorEntity[LidarrDataUpdateCoordinator[T]]): diff --git a/homeassistant/components/lidarr/config_flow.py b/homeassistant/components/lidarr/config_flow.py index 379a01375b6..05d6900bb41 100644 --- a/homeassistant/components/lidarr/config_flow.py +++ b/homeassistant/components/lidarr/config_flow.py @@ -10,11 +10,12 @@ from aiopyarr import exceptions from aiopyarr.lidarr_client import LidarrClient import voluptuous as vol -from homeassistant.config_entries import ConfigEntry, ConfigFlow, ConfigFlowResult +from homeassistant.config_entries import ConfigFlow, ConfigFlowResult from homeassistant.const import CONF_API_KEY, CONF_URL, CONF_VERIFY_SSL from homeassistant.core import HomeAssistant from homeassistant.helpers.aiohttp_client import async_get_clientsession +from . import LidarrConfigEntry from .const import DEFAULT_NAME, DOMAIN @@ -25,7 +26,7 @@ class LidarrConfigFlow(ConfigFlow, domain=DOMAIN): def __init__(self) -> None: """Initialize the flow.""" - self.entry: ConfigEntry | None = None + self.entry: LidarrConfigEntry | None = None async def async_step_reauth( self, user_input: Mapping[str, Any] diff --git a/homeassistant/components/lidarr/sensor.py b/homeassistant/components/lidarr/sensor.py index c876aec4623..b50a826a1c7 100644 --- a/homeassistant/components/lidarr/sensor.py +++ b/homeassistant/components/lidarr/sensor.py @@ -14,13 +14,12 @@ from homeassistant.components.sensor import ( SensorEntityDescription, SensorStateClass, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import UnitOfInformation from homeassistant.core import HomeAssistant from homeassistant.helpers.entity_platform import AddEntitiesCallback -from . import LidarrEntity -from .const import BYTE_SIZES, DOMAIN +from . import LidarrConfigEntry, LidarrEntity +from .const import BYTE_SIZES from .coordinator import LidarrDataUpdateCoordinator, T @@ -106,16 +105,13 @@ SENSOR_TYPES: dict[str, LidarrSensorEntityDescription[Any]] = { async def async_setup_entry( hass: HomeAssistant, - entry: ConfigEntry, + entry: LidarrConfigEntry, async_add_entities: AddEntitiesCallback, ) -> None: """Set up Lidarr sensors based on a config entry.""" - coordinators: dict[str, LidarrDataUpdateCoordinator[Any]] = hass.data[DOMAIN][ - entry.entry_id - ] entities: list[LidarrSensor[Any]] = [] for coordinator_type, description in SENSOR_TYPES.items(): - coordinator = coordinators[coordinator_type] + coordinator = getattr(entry.runtime_data, coordinator_type) if coordinator_type != "disk_space": entities.append(LidarrSensor(coordinator, description)) else: