evlli/hass-core
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Move Sonarr API calls to coordinators (#79826)

Browse source
This commit is contained in:
Robert Hillis 2022-10-07 18:25:16 -04:00 • committed by GitHub
parent 61901a1a60
commit 87a22fbcca
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 261 additions and 233 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

59
homeassistant/components/sonarr/__init__.py
View file

@ -1,10 +1,8 @@
"""The Sonarr component."""
from __future__ import annotations
from datetime import timedelta
import logging
from typing import Any
from aiopyarr import ArrAuthenticationException, ArrException
from aiopyarr.models.host_configuration import PyArrHostConfiguration
from aiopyarr.sonarr_client import SonarrClient
@ -19,24 +17,29 @@ from homeassistant.const import (
Platform,
)
from homeassistant.core import HomeAssistant
from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
from homeassistant.helpers.aiohttp_client import async_get_clientsession
from .const import (
CONF_BASE_PATH,
CONF_UPCOMING_DAYS,
CONF_WANTED_MAX_ITEMS,
DATA_HOST_CONFIG,
DATA_SONARR,
DATA_SYSTEM_STATUS,
DEFAULT_UPCOMING_DAYS,
DEFAULT_WANTED_MAX_ITEMS,
DOMAIN,
LOGGER,
)
from .coordinator import (
CalendarDataUpdateCoordinator,
CommandsDataUpdateCoordinator,
DiskSpaceDataUpdateCoordinator,
QueueDataUpdateCoordinator,
SeriesDataUpdateCoordinator,
SonarrDataUpdateCoordinator,
StatusDataUpdateCoordinator,
WantedDataUpdateCoordinator,
)
PLATFORMS = [Platform.SENSOR]
SCAN_INTERVAL = timedelta(seconds=30)
_LOGGER = logging.getLogger(__name__)
async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
@ -57,30 +60,28 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
url=entry.data[CONF_URL],
verify_ssl=entry.data[CONF_VERIFY_SSL],
)
sonarr = SonarrClient(
host_configuration=host_configuration,
session=async_get_clientsession(hass),
)
try:
system_status = await sonarr.async_get_system_status()
except ArrAuthenticationException as err:
raise ConfigEntryAuthFailed(
"API Key is no longer valid. Please reauthenticate"
) from err
except ArrException as err:
raise ConfigEntryNotReady from err
entry.async_on_unload(entry.add_update_listener(_async_update_listener))
hass.data.setdefault(DOMAIN, {})
hass.data[DOMAIN][entry.entry_id] = {
DATA_HOST_CONFIG: host_configuration,
DATA_SONARR: sonarr,
DATA_SYSTEM_STATUS: system_status,
coordinators: dict[str, SonarrDataUpdateCoordinator[Any]] = {
"upcoming": CalendarDataUpdateCoordinator(hass, host_configuration, sonarr),
"commands": CommandsDataUpdateCoordinator(hass, host_configuration, sonarr),
"diskspace": DiskSpaceDataUpdateCoordinator(hass, host_configuration, sonarr),
"queue": QueueDataUpdateCoordinator(hass, host_configuration, sonarr),
"series": SeriesDataUpdateCoordinator(hass, host_configuration, sonarr),
"status": StatusDataUpdateCoordinator(hass, host_configuration, sonarr),
"wanted": WantedDataUpdateCoordinator(hass, host_configuration, sonarr),
}
# Temporary, until we add diagnostic entities
_version = None
for coordinator in coordinators.values():
await coordinator.async_config_entry_first_refresh()
if isinstance(coordinator, StatusDataUpdateCoordinator):
_version = coordinator.data.version
coordinator.system_version = _version
hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinators
await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
return True
@ -88,7 +89,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
async def async_migrate_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Migrate old entry."""
_LOGGER.debug("Migrating from version %s", entry.version)
LOGGER.debug("Migrating from version %s", entry.version)
if entry.version == 1:
new_proto = "https" if entry.data[CONF_SSL] else "http"
@ -106,7 +107,7 @@ async def async_migrate_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
hass.config_entries.async_update_entry(entry, data=data)
entry.version = 2
_LOGGER.info("Migration to version %s successful", entry.version)
LOGGER.info("Migration to version %s successful", entry.version)
return True

9
homeassistant/components/sonarr/const.py
View file

@ -1,4 +1,6 @@
"""Constants for Sonarr."""
import logging
DOMAIN = "sonarr"
# Config Keys
@ -9,12 +11,9 @@ CONF_UNIT = "unit"
CONF_UPCOMING_DAYS = "upcoming_days"
CONF_WANTED_MAX_ITEMS = "wanted_max_items"
# Data
DATA_HOST_CONFIG = "host_config"
DATA_SONARR = "sonarr"
DATA_SYSTEM_STATUS = "system_status"
# Defaults
DEFAULT_UPCOMING_DAYS = 1
DEFAULT_VERIFY_SSL = False
DEFAULT_WANTED_MAX_ITEMS = 50
LOGGER = logging.getLogger(__package__)

147
homeassistant/components/sonarr/coordinator.py Normal file
View file

@ -0,0 +1,147 @@
"""Data update coordinator for the Sonarr integration."""
from __future__ import annotations
from datetime import timedelta
from typing import TypeVar, Union, cast
from aiopyarr import (
Command,
Diskspace,
SonarrCalendar,
SonarrQueue,
SonarrSeries,
SonarrWantedMissing,
SystemStatus,
exceptions,
)
from aiopyarr.models.host_configuration import PyArrHostConfiguration
from aiopyarr.sonarr_client import SonarrClient
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant
from homeassistant.exceptions import ConfigEntryAuthFailed
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
import homeassistant.util.dt as dt_util
from .const import CONF_UPCOMING_DAYS, CONF_WANTED_MAX_ITEMS, DOMAIN, LOGGER
SonarrDataT = TypeVar(
"SonarrDataT",
bound=Union[
list[SonarrCalendar],
list[Command],
list[Diskspace],
SonarrQueue,
list[SonarrSeries],
SystemStatus,
SonarrWantedMissing,
],
)
class SonarrDataUpdateCoordinator(DataUpdateCoordinator[SonarrDataT]):
"""Data update coordinator for the Sonarr integration."""
config_entry: ConfigEntry
def __init__(
self,
hass: HomeAssistant,
host_configuration: PyArrHostConfiguration,
api_client: SonarrClient,
) -> None:
"""Initialize the coordinator."""
super().__init__(
hass=hass,
logger=LOGGER,
name=DOMAIN,
update_interval=timedelta(seconds=30),
)
self.api_client = api_client
self.host_configuration = host_configuration
self.system_version: str | None = None
async def _async_update_data(self) -> SonarrDataT:
"""Get the latest data from Sonarr."""
try:
return await self._fetch_data()
except exceptions.ArrConnectionException as ex:
raise UpdateFailed(ex) from ex
except exceptions.ArrAuthenticationException as ex:
raise ConfigEntryAuthFailed(
"API Key is no longer valid. Please reauthenticate"
) from ex
async def _fetch_data(self) -> SonarrDataT:
"""Fetch the actual data."""
raise NotImplementedError
class CalendarDataUpdateCoordinator(SonarrDataUpdateCoordinator[list[SonarrCalendar]]):
"""Calendar update coordinator."""
async def _fetch_data(self) -> list[SonarrCalendar]:
"""Fetch the movies data."""
local = dt_util.start_of_local_day().replace(microsecond=0)
start = dt_util.as_utc(local)
end = start + timedelta(days=self.config_entry.options[CONF_UPCOMING_DAYS])
return cast(
list[SonarrCalendar],
await self.api_client.async_get_calendar(
start_date=start, end_date=end, include_series=True
),
)
class CommandsDataUpdateCoordinator(SonarrDataUpdateCoordinator[list[Command]]):
"""Commands update coordinator for Sonarr."""
async def _fetch_data(self) -> list[Command]:
"""Fetch the data."""
return cast(list[Command], await self.api_client.async_get_commands())
class DiskSpaceDataUpdateCoordinator(SonarrDataUpdateCoordinator[list[Diskspace]]):
"""Disk space update coordinator for Sonarr."""
async def _fetch_data(self) -> list[Diskspace]:
"""Fetch the data."""
return await self.api_client.async_get_diskspace()
class QueueDataUpdateCoordinator(SonarrDataUpdateCoordinator[SonarrQueue]):
"""Queue update coordinator."""
async def _fetch_data(self) -> SonarrQueue:
"""Fetch the data."""
return await self.api_client.async_get_queue(
include_series=True, include_episode=True
)
class SeriesDataUpdateCoordinator(SonarrDataUpdateCoordinator[list[SonarrSeries]]):
"""Series update coordinator."""
async def _fetch_data(self) -> list[SonarrSeries]:
"""Fetch the data."""
return cast(list[SonarrSeries], await self.api_client.async_get_series())
class StatusDataUpdateCoordinator(SonarrDataUpdateCoordinator[SystemStatus]):
"""Status update coordinator for Sonarr."""
async def _fetch_data(self) -> SystemStatus:
"""Fetch the data."""
return await self.api_client.async_get_system_status()
class WantedDataUpdateCoordinator(SonarrDataUpdateCoordinator[SonarrWantedMissing]):
"""Wanted update coordinator."""
async def _fetch_data(self) -> SonarrWantedMissing:
"""Fetch the data."""
return await self.api_client.async_get_wanted(
page_size=self.config_entry.options[CONF_WANTED_MAX_ITEMS],
include_series=True,
)

37
homeassistant/components/sonarr/entity.py
View file

@ -1,41 +1,36 @@
"""Base Entity for Sonarr."""
from aiopyarr import SystemStatus
from aiopyarr.models.host_configuration import PyArrHostConfiguration
from aiopyarr.sonarr_client import SonarrClient
from __future__ import annotations
from homeassistant.helpers.device_registry import DeviceEntryType
from homeassistant.helpers.entity import DeviceInfo, Entity
from homeassistant.helpers.entity import DeviceInfo, EntityDescription
from homeassistant.helpers.update_coordinator import CoordinatorEntity
from .const import DOMAIN
from .coordinator import SonarrDataT, SonarrDataUpdateCoordinator
class SonarrEntity(Entity):
class SonarrEntity(CoordinatorEntity[SonarrDataUpdateCoordinator[SonarrDataT]]):
"""Defines a base Sonarr entity."""
def __init__(
self,
*,
sonarr: SonarrClient,
host_config: PyArrHostConfiguration,
system_status: SystemStatus,
entry_id: str,
device_id: str,
coordinator: SonarrDataUpdateCoordinator[SonarrDataT],
description: EntityDescription,
) -> None:
"""Initialize the Sonarr entity."""
self._entry_id = entry_id
self._device_id = device_id
self.sonarr = sonarr
self.host_config = host_config
self.system_status = system_status
super().__init__(coordinator)
self.coordinator = coordinator
self.entity_description = description
self._attr_unique_id = f"{coordinator.config_entry.entry_id}_{description.key}"
@property
def device_info(self) -> DeviceInfo:
"""Return device information about the application."""
return DeviceInfo(
identifiers={(DOMAIN, self._device_id)},
name="Activity Sensor",
manufacturer="Sonarr",
sw_version=self.system_status.version,
configuration_url=self.coordinator.host_configuration.base_url,
entry_type=DeviceEntryType.SERVICE,
configuration_url=self.host_config.base_url,
identifiers={(DOMAIN, self.coordinator.config_entry.entry_id)},
manufacturer="Sonarr",
name="Activity Sensor",
sw_version=self.coordinator.system_version,
)

241
homeassistant/components/sonarr/sensor.py
View file

@ -1,16 +1,17 @@
"""Support for Sonarr sensors."""
from __future__ import annotations
from collections.abc import Awaitable, Callable, Coroutine
from datetime import timedelta
from functools import wraps
import logging
from typing import Any, TypeVar, cast
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any, Generic
from aiopyarr import ArrConnectionException, ArrException, SystemStatus
from aiopyarr.models.host_configuration import PyArrHostConfiguration
from aiopyarr.sonarr_client import SonarrClient
from typing_extensions import Concatenate, ParamSpec
from aiopyarr import (
Diskspace,
SonarrCalendar,
SonarrQueue,
SonarrSeries,
SonarrWantedMissing,
)
from homeassistant.components.sensor import SensorEntity, SensorEntityDescription
from homeassistant.config_entries import ConfigEntry
@ -20,64 +21,74 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.typing import StateType
import homeassistant.util.dt as dt_util
from .const import (
CONF_UPCOMING_DAYS,
CONF_WANTED_MAX_ITEMS,
DATA_HOST_CONFIG,
DATA_SONARR,
DATA_SYSTEM_STATUS,
DOMAIN,
)
from .const import DOMAIN
from .coordinator import SonarrDataT, SonarrDataUpdateCoordinator
from .entity import SonarrEntity
_LOGGER = logging.getLogger(__name__)
SENSOR_TYPES: tuple[SensorEntityDescription, ...] = (
SensorEntityDescription(
@dataclass
class SonarrSensorEntityDescriptionMixIn(Generic[SonarrDataT]):
"""Mixin for Sonarr sensor."""
value_fn: Callable[[SonarrDataT], StateType]
@dataclass
class SonarrSensorEntityDescription(
SensorEntityDescription, SonarrSensorEntityDescriptionMixIn[SonarrDataT]
):
"""Class to describe a Sonarr sensor."""
SENSOR_TYPES: dict[str, SonarrSensorEntityDescription[Any]] = {
"commands": SonarrSensorEntityDescription(
key="commands",
name="Sonarr Commands",
icon="mdi:code-braces",
native_unit_of_measurement="Commands",
entity_registry_enabled_default=False,
value_fn=len,
),
SensorEntityDescription(
"diskspace": SonarrSensorEntityDescription[list[Diskspace]](
key="diskspace",
name="Sonarr Disk Space",
icon="mdi:harddisk",
native_unit_of_measurement=DATA_GIGABYTES,
entity_registry_enabled_default=False,
value_fn=lambda data: f"{sum(disk.freeSpace for disk in data) / 1024**3:.2f}",
),
SensorEntityDescription(
"queue": SonarrSensorEntityDescription[SonarrQueue](
key="queue",
name="Sonarr Queue",
icon="mdi:download",
native_unit_of_measurement="Episodes",
entity_registry_enabled_default=False,
value_fn=lambda data: data.totalRecords,
),
SensorEntityDescription(
"series": SonarrSensorEntityDescription[list[SonarrSeries]](
key="series",
name="Sonarr Shows",
icon="mdi:television",
native_unit_of_measurement="Series",
entity_registry_enabled_default=False,
value_fn=len,
),
SensorEntityDescription(
"upcoming": SonarrSensorEntityDescription[list[SonarrCalendar]](
key="upcoming",
name="Sonarr Upcoming",
icon="mdi:television",
native_unit_of_measurement="Episodes",
value_fn=len,
),
SensorEntityDescription(
"wanted": SonarrSensorEntityDescription[SonarrWantedMissing](
key="wanted",
name="Sonarr Wanted",
icon="mdi:television",
native_unit_of_measurement="Episodes",
entity_registry_enabled_default=False,
value_fn=lambda data: data.totalRecords,
),
)
_SonarrSensorT = TypeVar("_SonarrSensorT", bound="SonarrSensor")
_P = ParamSpec("_P")
}
async def async_setup_entry(
@ -86,134 +97,30 @@ async def async_setup_entry(
async_add_entities: AddEntitiesCallback,
) -> None:
"""Set up Sonarr sensors based on a config entry."""
sonarr: SonarrClient = hass.data[DOMAIN][entry.entry_id][DATA_SONARR]
host_config: PyArrHostConfiguration = hass.data[DOMAIN][entry.entry_id][
DATA_HOST_CONFIG
coordinators: dict[str, SonarrDataUpdateCoordinator[Any]] = hass.data[DOMAIN][
entry.entry_id
]
system_status: SystemStatus = hass.data[DOMAIN][entry.entry_id][DATA_SYSTEM_STATUS]
options: dict[str, Any] = dict(entry.options)
entities = [
SonarrSensor(
sonarr,
host_config,
system_status,
entry.entry_id,
description,
options,
)
for description in SENSOR_TYPES
]
async_add_entities(entities, True)
async_add_entities(
SonarrSensor(coordinators[coordinator_type], description)
for coordinator_type, description in SENSOR_TYPES.items()
)
def sonarr_exception_handler(
func: Callable[Concatenate[_SonarrSensorT, _P], Awaitable[None]]
) -> Callable[Concatenate[_SonarrSensorT, _P], Coroutine[Any, Any, None]]:
"""Decorate Sonarr calls to handle Sonarr exceptions.
A decorator that wraps the passed in function, catches Sonarr errors,
and handles the availability of the entity.
"""
@wraps(func)
async def wrapper(
self: _SonarrSensorT, *args: _P.args, **kwargs: _P.kwargs
) -> None:
try:
await func(self, *args, **kwargs)
self.last_update_success = True
except ArrConnectionException as error:
if self.last_update_success:
_LOGGER.error("Error communicating with API: %s", error)
self.last_update_success = False
except ArrException as error:
if self.last_update_success:
_LOGGER.error("Invalid response from API: %s", error)
self.last_update_success = False
return wrapper
class SonarrSensor(SonarrEntity, SensorEntity):
class SonarrSensor(SonarrEntity[SonarrDataT], SensorEntity):
"""Implementation of the Sonarr sensor."""
data: dict[str, Any]
last_update_success: bool
upcoming_days: int
wanted_max_items: int
def __init__(
self,
sonarr: SonarrClient,
host_config: PyArrHostConfiguration,
system_status: SystemStatus,
entry_id: str,
description: SensorEntityDescription,
options: dict[str, Any],
) -> None:
"""Initialize Sonarr sensor."""
self.entity_description = description
self._attr_unique_id = f"{entry_id}_{description.key}"
self.data = {}
self.last_update_success = True
self.upcoming_days = options[CONF_UPCOMING_DAYS]
self.wanted_max_items = options[CONF_WANTED_MAX_ITEMS]
super().__init__(
sonarr=sonarr,
host_config=host_config,
system_status=system_status,
entry_id=entry_id,
device_id=entry_id,
)
@property
def available(self) -> bool:
"""Return sensor availability."""
return self.last_update_success
@sonarr_exception_handler
async def async_update(self) -> None:
"""Update entity."""
key = self.entity_description.key
if key == "diskspace":
self.data[key] = await self.sonarr.async_get_diskspace()
elif key == "commands":
self.data[key] = await self.sonarr.async_get_commands()
elif key == "queue":
self.data[key] = await self.sonarr.async_get_queue(
include_series=True, include_episode=True
)
elif key == "series":
self.data[key] = await self.sonarr.async_get_series()
elif key == "upcoming":
local = dt_util.start_of_local_day().replace(microsecond=0)
start = dt_util.as_utc(local)
end = start + timedelta(days=self.upcoming_days)
self.data[key] = await self.sonarr.async_get_calendar(
start_date=start,
end_date=end,
include_series=True,
)
elif key == "wanted":
self.data[key] = await self.sonarr.async_get_wanted(
page_size=self.wanted_max_items,
include_series=True,
)
coordinator: SonarrDataUpdateCoordinator
entity_description: SonarrSensorEntityDescription[SonarrDataT]
@property
def extra_state_attributes(self) -> dict[str, str] | None:
"""Return the state attributes of the entity."""
attrs = {}
key = self.entity_description.key
data = self.coordinator.data
if key == "diskspace" and self.data.get(key) is not None:
for disk in self.data[key]:
if key == "diskspace":
for disk in data:
free = disk.freeSpace / 1024**3
total = disk.totalSpace / 1024**3
usage = free / total * 100
@ -221,29 +128,29 @@ class SonarrSensor(SonarrEntity, SensorEntity):
attrs[
disk.path
] = f"{free:.2f}/{total:.2f}{self.unit_of_measurement} ({usage:.2f}%)"
elif key == "commands" and self.data.get(key) is not None:
for command in self.data[key]:
elif key == "commands":
for command in data:
attrs[command.name] = command.status
elif key == "queue" and self.data.get(key) is not None:
for item in self.data[key].records:
elif key == "queue":
for item in data.records:
remaining = 1 if item.size == 0 else item.sizeleft / item.size
remaining_pct = 100 * (1 - remaining)
identifier = f"S{item.episode.seasonNumber:02d}E{item.episode. episodeNumber:02d}"
name = f"{item.series.title} {identifier}"
attrs[name] = f"{remaining_pct:.2f}%"
elif key == "series" and self.data.get(key) is not None:
for item in self.data[key]:
elif key == "series":
for item in data:
stats = item.statistics
attrs[
item.title
] = f"{getattr(stats,'episodeFileCount', 0)}/{getattr(stats, 'episodeCount', 0)} Episodes"
elif key == "upcoming" and self.data.get(key) is not None:
for episode in self.data[key]:
elif key == "upcoming":
for episode in data:
identifier = f"S{episode.seasonNumber:02d}E{episode.episodeNumber:02d}"
attrs[episode.series.title] = identifier
elif key == "wanted" and self.data.get(key) is not None:
for item in self.data[key].records:
elif key == "wanted":
for item in data.records:
identifier = f"S{item.seasonNumber:02d}E{item.episodeNumber:02d}"
name = f"{item.series.title} {identifier}"
@ -256,26 +163,4 @@ class SonarrSensor(SonarrEntity, SensorEntity):
@property
def native_value(self) -> StateType:
"""Return the state of the sensor."""
key = self.entity_description.key
if key == "diskspace" and self.data.get(key) is not None:
total_free = sum(disk.freeSpace for disk in self.data[key])
free = total_free / 1024**3
return f"{free:.2f}"
if key == "commands" and self.data.get(key) is not None:
return len(self.data[key])
if key == "queue" and self.data.get(key) is not None:
return cast(int, self.data[key].totalRecords)
if key == "series" and self.data.get(key) is not None:
return len(self.data[key])
if key == "upcoming" and self.data.get(key) is not None:
return len(self.data[key])
if key == "wanted" and self.data.get(key) is not None:
return cast(int, self.data[key].totalRecords)
return None
return self.entity_description.value_fn(self.coordinator.data)

1
tests/components/sonarr/fixtures/system-status.json
View file

@ -1,5 +1,6 @@
{
"appName": "Sonarr",
"instanceName": "Sonarr",
"version": "3.0.6.1451",
"buildTime": "2022-01-23T16:51:56Z",
"isDebug": false,