Store runtime data inside the config entry in Lidarr (#119548)

This commit is contained in:
Robert Hillis 2024-06-18 03:13:24 -04:00 committed by GitHub
parent d5d906e148
commit 67223b2a2d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 32 additions and 25 deletions

View File

@ -2,7 +2,7 @@
from __future__ import annotations from __future__ import annotations
from typing import Any from dataclasses import dataclass, fields
from aiopyarr.lidarr_client import LidarrClient from aiopyarr.lidarr_client import LidarrClient
from aiopyarr.models.host_configuration import PyArrHostConfiguration from aiopyarr.models.host_configuration import PyArrHostConfiguration
@ -25,10 +25,22 @@ from .coordinator import (
WantedDataUpdateCoordinator, WantedDataUpdateCoordinator,
) )
type LidarrConfigEntry = ConfigEntry[LidarrData]
PLATFORMS = [Platform.SENSOR] 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.""" """Set up Lidarr from a config entry."""
host_configuration = PyArrHostConfiguration( host_configuration = PyArrHostConfiguration(
api_token=entry.data[CONF_API_KEY], 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), session=async_get_clientsession(hass, host_configuration.verify_ssl),
request_timeout=60, request_timeout=60,
) )
coordinators: dict[str, LidarrDataUpdateCoordinator[Any]] = { data = LidarrData(
"disk_space": DiskSpaceDataUpdateCoordinator(hass, host_configuration, lidarr), disk_space=DiskSpaceDataUpdateCoordinator(hass, host_configuration, lidarr),
"queue": QueueDataUpdateCoordinator(hass, host_configuration, lidarr), queue=QueueDataUpdateCoordinator(hass, host_configuration, lidarr),
"status": StatusDataUpdateCoordinator(hass, host_configuration, lidarr), status=StatusDataUpdateCoordinator(hass, host_configuration, lidarr),
"wanted": WantedDataUpdateCoordinator(hass, host_configuration, lidarr), wanted=WantedDataUpdateCoordinator(hass, host_configuration, lidarr),
} )
# Temporary, until we add diagnostic entities # Temporary, until we add diagnostic entities
_version = None _version = None
for coordinator in coordinators.values(): for field in fields(data):
coordinator = getattr(data, field.name)
await coordinator.async_config_entry_first_refresh() await coordinator.async_config_entry_first_refresh()
if isinstance(coordinator, StatusDataUpdateCoordinator): if isinstance(coordinator, StatusDataUpdateCoordinator):
_version = coordinator.data _version = coordinator.data
coordinator.system_version = _version 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) await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
return True 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.""" """Unload a 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
class LidarrEntity(CoordinatorEntity[LidarrDataUpdateCoordinator[T]]): class LidarrEntity(CoordinatorEntity[LidarrDataUpdateCoordinator[T]]):

View File

@ -10,11 +10,12 @@ from aiopyarr import exceptions
from aiopyarr.lidarr_client import LidarrClient from aiopyarr.lidarr_client import LidarrClient
import voluptuous as vol 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.const import CONF_API_KEY, CONF_URL, CONF_VERIFY_SSL
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from homeassistant.helpers.aiohttp_client import async_get_clientsession from homeassistant.helpers.aiohttp_client import async_get_clientsession
from . import LidarrConfigEntry
from .const import DEFAULT_NAME, DOMAIN from .const import DEFAULT_NAME, DOMAIN
@ -25,7 +26,7 @@ class LidarrConfigFlow(ConfigFlow, domain=DOMAIN):
def __init__(self) -> None: def __init__(self) -> None:
"""Initialize the flow.""" """Initialize the flow."""
self.entry: ConfigEntry | None = None self.entry: LidarrConfigEntry | None = None
async def async_step_reauth( async def async_step_reauth(
self, user_input: Mapping[str, Any] self, user_input: Mapping[str, Any]

View File

@ -14,13 +14,12 @@ from homeassistant.components.sensor import (
SensorEntityDescription, SensorEntityDescription,
SensorStateClass, SensorStateClass,
) )
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import UnitOfInformation from homeassistant.const import UnitOfInformation
from homeassistant.core import HomeAssistant from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_platform import AddEntitiesCallback
from . import LidarrEntity from . import LidarrConfigEntry, LidarrEntity
from .const import BYTE_SIZES, DOMAIN from .const import BYTE_SIZES
from .coordinator import LidarrDataUpdateCoordinator, T from .coordinator import LidarrDataUpdateCoordinator, T
@ -106,16 +105,13 @@ SENSOR_TYPES: dict[str, LidarrSensorEntityDescription[Any]] = {
async def async_setup_entry( async def async_setup_entry(
hass: HomeAssistant, hass: HomeAssistant,
entry: ConfigEntry, entry: LidarrConfigEntry,
async_add_entities: AddEntitiesCallback, async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up Lidarr sensors based on a config entry.""" """Set up Lidarr sensors based on a config entry."""
coordinators: dict[str, LidarrDataUpdateCoordinator[Any]] = hass.data[DOMAIN][
entry.entry_id
]
entities: list[LidarrSensor[Any]] = [] entities: list[LidarrSensor[Any]] = []
for coordinator_type, description in SENSOR_TYPES.items(): for coordinator_type, description in SENSOR_TYPES.items():
coordinator = coordinators[coordinator_type] coordinator = getattr(entry.runtime_data, coordinator_type)
if coordinator_type != "disk_space": if coordinator_type != "disk_space":
entities.append(LidarrSensor(coordinator, description)) entities.append(LidarrSensor(coordinator, description))
else: else: