Migrate sabnzbd to use data update coordinator (#114745)

* Migrate sabnzbd to use data update coordinator

* Add to coveragerc

* Setup coordinator after migration

* Use kB/s as UoM

* Add suggested
This commit is contained in:
Jan-Philipp Benecke 2024-04-03 15:15:23 +02:00 committed by GitHub
parent 2b9f22f11e
commit 613bdebfe5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 88 additions and 124 deletions

View file

@ -1184,6 +1184,7 @@ omit =
homeassistant/components/rympro/coordinator.py
homeassistant/components/rympro/sensor.py
homeassistant/components/sabnzbd/__init__.py
homeassistant/components/sabnzbd/coordinator.py
homeassistant/components/sabnzbd/sensor.py
homeassistant/components/saj/sensor.py
homeassistant/components/satel_integra/*

View file

@ -1177,8 +1177,8 @@ build.json @home-assistant/supervisor
/tests/components/ruuvitag_ble/ @akx
/homeassistant/components/rympro/ @OnFreund @elad-bar @maorcc
/tests/components/rympro/ @OnFreund @elad-bar @maorcc
/homeassistant/components/sabnzbd/ @shaiu
/tests/components/sabnzbd/ @shaiu
/homeassistant/components/sabnzbd/ @shaiu @jpbede
/tests/components/sabnzbd/ @shaiu @jpbede
/homeassistant/components/saj/ @fredericvl
/homeassistant/components/samsungtv/ @chemelli74 @epenet
/tests/components/samsungtv/ @chemelli74 @epenet

View file

@ -6,7 +6,6 @@ from collections.abc import Callable, Coroutine
import logging
from typing import Any
from pysabnzbd import SabnzbdApiException
import voluptuous as vol
from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry, ConfigEntryState
@ -23,9 +22,7 @@ from homeassistant.core import HomeAssistant, ServiceCall, callback
from homeassistant.exceptions import ConfigEntryNotReady, HomeAssistantError
import homeassistant.helpers.config_validation as cv
from homeassistant.helpers.device_registry import async_get
from homeassistant.helpers.dispatcher import async_dispatcher_send
from homeassistant.helpers.entity_registry import RegistryEntry, async_migrate_entries
from homeassistant.helpers.event import async_track_time_interval
from homeassistant.helpers.typing import ConfigType
from .const import (
@ -37,15 +34,11 @@ from .const import (
DEFAULT_SPEED_LIMIT,
DEFAULT_SSL,
DOMAIN,
KEY_API,
KEY_API_DATA,
KEY_NAME,
SERVICE_PAUSE,
SERVICE_RESUME,
SERVICE_SET_SPEED,
SIGNAL_SABNZBD_UPDATED,
UPDATE_INTERVAL,
)
from .coordinator import SabnzbdUpdateCoordinator
from .sab import get_client
from .sensor import OLD_SENSOR_KEYS
@ -179,30 +172,28 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
if not sab_api:
raise ConfigEntryNotReady
sab_api_data = SabnzbdApiData(sab_api)
hass.data.setdefault(DOMAIN, {})[entry.entry_id] = {
KEY_API: sab_api,
KEY_API_DATA: sab_api_data,
KEY_NAME: entry.data[CONF_NAME],
}
await migrate_unique_id(hass, entry)
update_device_identifiers(hass, entry)
coordinator = SabnzbdUpdateCoordinator(hass, sab_api)
await coordinator.async_config_entry_first_refresh()
hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
@callback
def extract_api(
func: Callable[[ServiceCall, SabnzbdApiData], Coroutine[Any, Any, None]],
func: Callable[
[ServiceCall, SabnzbdUpdateCoordinator], Coroutine[Any, Any, None]
],
) -> Callable[[ServiceCall], Coroutine[Any, Any, None]]:
"""Define a decorator to get the correct api for a service call."""
async def wrapper(call: ServiceCall) -> None:
"""Wrap the service function."""
entry_id = async_get_entry_id_for_service_call(hass, call)
api_data = hass.data[DOMAIN][entry_id][KEY_API_DATA]
coordinator: SabnzbdUpdateCoordinator = hass.data[DOMAIN][entry_id]
try:
await func(call, api_data)
await func(call, coordinator)
except Exception as err:
raise HomeAssistantError(
f"Error while executing {func.__name__}: {err}"
@ -211,17 +202,23 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
return wrapper
@extract_api
async def async_pause_queue(call: ServiceCall, api: SabnzbdApiData) -> None:
await api.async_pause_queue()
async def async_pause_queue(
call: ServiceCall, coordinator: SabnzbdUpdateCoordinator
) -> None:
await coordinator.sab_api.pause_queue()
@extract_api
async def async_resume_queue(call: ServiceCall, api: SabnzbdApiData) -> None:
await api.async_resume_queue()
async def async_resume_queue(
call: ServiceCall, coordinator: SabnzbdUpdateCoordinator
) -> None:
await coordinator.sab_api.resume_queue()
@extract_api
async def async_set_queue_speed(call: ServiceCall, api: SabnzbdApiData) -> None:
async def async_set_queue_speed(
call: ServiceCall, coordinator: SabnzbdUpdateCoordinator
) -> None:
speed = call.data.get(ATTR_SPEED)
await api.async_set_queue_speed(speed)
await coordinator.sab_api.set_speed_limit(speed)
for service, method, schema in (
(SERVICE_PAUSE, async_pause_queue, SERVICE_BASE_SCHEMA),
@ -233,18 +230,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
hass.services.async_register(DOMAIN, service, method, schema=schema)
async def async_update_sabnzbd(now):
"""Refresh SABnzbd queue data."""
try:
await sab_api.refresh_data()
async_dispatcher_send(hass, SIGNAL_SABNZBD_UPDATED, None)
except SabnzbdApiException as err:
_LOGGER.error(err)
entry.async_on_unload(
async_track_time_interval(hass, async_update_sabnzbd, UPDATE_INTERVAL)
)
await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
return True
@ -268,42 +253,3 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
hass.services.async_remove(DOMAIN, service_name)
return unload_ok
class SabnzbdApiData:
"""Class for storing/refreshing sabnzbd api queue data."""
def __init__(self, sab_api):
"""Initialize component."""
self.sab_api = sab_api
async def async_pause_queue(self):
"""Pause Sabnzbd queue."""
try:
return await self.sab_api.pause_queue()
except SabnzbdApiException as err:
_LOGGER.error(err)
return False
async def async_resume_queue(self):
"""Resume Sabnzbd queue."""
try:
return await self.sab_api.resume_queue()
except SabnzbdApiException as err:
_LOGGER.error(err)
return False
async def async_set_queue_speed(self, limit):
"""Set speed limit for the Sabnzbd queue."""
try:
return await self.sab_api.set_speed_limit(limit)
except SabnzbdApiException as err:
_LOGGER.error(err)
return False
def get_queue_field(self, field):
"""Return the value for the given field from the Sabnzbd queue."""
return self.sab_api.queue.get(field)

View file

@ -1,7 +1,5 @@
"""Constants for the Sabnzbd component."""
from datetime import timedelta
DOMAIN = "sabnzbd"
DATA_SABNZBD = "sabnzbd"
@ -14,14 +12,6 @@ DEFAULT_PORT = 8080
DEFAULT_SPEED_LIMIT = "100"
DEFAULT_SSL = False
UPDATE_INTERVAL = timedelta(seconds=30)
SERVICE_PAUSE = "pause"
SERVICE_RESUME = "resume"
SERVICE_SET_SPEED = "set_speed"
SIGNAL_SABNZBD_UPDATED = "sabnzbd_updated"
KEY_API = "api"
KEY_API_DATA = "api_data"
KEY_NAME = "name"

View file

@ -0,0 +1,40 @@
"""DataUpdateCoordinator for the SABnzbd integration."""
from datetime import timedelta
import logging
from typing import Any
from pysabnzbd import SabnzbdApi, SabnzbdApiException
from homeassistant.core import HomeAssistant
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
_LOGGER = logging.getLogger(__name__)
class SabnzbdUpdateCoordinator(DataUpdateCoordinator[dict[str, Any]]):
"""The SABnzbd update coordinator."""
def __init__(
self,
hass: HomeAssistant,
sab_api: SabnzbdApi,
) -> None:
"""Initialize the SABnzbd update coordinator."""
self.sab_api = sab_api
super().__init__(
hass,
_LOGGER,
name="SABnzbd",
update_interval=timedelta(seconds=30),
)
async def _async_update_data(self) -> dict[str, Any]:
"""Get the latest data from the SABnzbd API."""
try:
await self.sab_api.refresh_data()
except SabnzbdApiException as err:
raise UpdateFailed("Error while fetching data") from err
return self.sab_api.queue

View file

@ -1,7 +1,7 @@
{
"domain": "sabnzbd",
"name": "SABnzbd",
"codeowners": ["@shaiu"],
"codeowners": ["@shaiu", "@jpbede"],
"config_flow": true,
"documentation": "https://www.home-assistant.io/integrations/sabnzbd",
"iot_class": "local_polling",

View file

@ -14,11 +14,12 @@ from homeassistant.config_entries import ConfigEntry
from homeassistant.const import UnitOfDataRate, UnitOfInformation
from homeassistant.core import HomeAssistant
from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
from homeassistant.helpers.dispatcher import async_dispatcher_connect
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.typing import StateType
from homeassistant.helpers.update_coordinator import CoordinatorEntity
from . import DOMAIN, SIGNAL_SABNZBD_UPDATED
from .const import DEFAULT_NAME, KEY_API_DATA
from . import DOMAIN, SabnzbdUpdateCoordinator
from .const import DEFAULT_NAME
@dataclass(frozen=True, kw_only=True)
@ -28,18 +29,18 @@ class SabnzbdSensorEntityDescription(SensorEntityDescription):
key: str
SPEED_KEY = "kbpersec"
SENSOR_TYPES: tuple[SabnzbdSensorEntityDescription, ...] = (
SabnzbdSensorEntityDescription(
key="status",
translation_key="status",
),
SabnzbdSensorEntityDescription(
key=SPEED_KEY,
key="kbpersec",
translation_key="speed",
device_class=SensorDeviceClass.DATA_RATE,
native_unit_of_measurement=UnitOfDataRate.MEGABYTES_PER_SECOND,
native_unit_of_measurement=UnitOfDataRate.KILOBYTES_PER_SECOND,
suggested_unit_of_measurement=UnitOfDataRate.MEGABYTES_PER_SECOND,
suggested_display_precision=1,
state_class=SensorStateClass.MEASUREMENT,
),
SabnzbdSensorEntityDescription(
@ -74,6 +75,7 @@ SENSOR_TYPES: tuple[SabnzbdSensorEntityDescription, ...] = (
key="noofslots_total",
translation_key="queue_count",
state_class=SensorStateClass.TOTAL,
suggested_display_precision=2,
),
SabnzbdSensorEntityDescription(
key="day_size",
@ -82,6 +84,7 @@ SENSOR_TYPES: tuple[SabnzbdSensorEntityDescription, ...] = (
device_class=SensorDeviceClass.DATA_SIZE,
entity_registry_enabled_default=False,
state_class=SensorStateClass.TOTAL_INCREASING,
suggested_display_precision=2,
),
SabnzbdSensorEntityDescription(
key="week_size",
@ -90,6 +93,7 @@ SENSOR_TYPES: tuple[SabnzbdSensorEntityDescription, ...] = (
device_class=SensorDeviceClass.DATA_SIZE,
entity_registry_enabled_default=False,
state_class=SensorStateClass.TOTAL_INCREASING,
suggested_display_precision=2,
),
SabnzbdSensorEntityDescription(
key="month_size",
@ -98,6 +102,7 @@ SENSOR_TYPES: tuple[SabnzbdSensorEntityDescription, ...] = (
device_class=SensorDeviceClass.DATA_SIZE,
entity_registry_enabled_default=False,
state_class=SensorStateClass.TOTAL_INCREASING,
suggested_display_precision=2,
),
SabnzbdSensorEntityDescription(
key="total_size",
@ -105,6 +110,7 @@ SENSOR_TYPES: tuple[SabnzbdSensorEntityDescription, ...] = (
native_unit_of_measurement=UnitOfInformation.GIGABYTES,
device_class=SensorDeviceClass.DATA_SIZE,
state_class=SensorStateClass.TOTAL_INCREASING,
suggested_display_precision=2,
),
)
@ -131,15 +137,14 @@ async def async_setup_entry(
"""Set up a Sabnzbd sensor entry."""
entry_id = config_entry.entry_id
sab_api_data = hass.data[DOMAIN][entry_id][KEY_API_DATA]
coordinator: SabnzbdUpdateCoordinator = hass.data[DOMAIN][entry_id]
async_add_entities(
[SabnzbdSensor(sab_api_data, sensor, entry_id) for sensor in SENSOR_TYPES]
[SabnzbdSensor(coordinator, sensor, entry_id) for sensor in SENSOR_TYPES]
)
class SabnzbdSensor(SensorEntity):
class SabnzbdSensor(CoordinatorEntity[SabnzbdUpdateCoordinator], SensorEntity):
"""Representation of an SABnzbd sensor."""
entity_description: SabnzbdSensorEntityDescription
@ -148,40 +153,22 @@ class SabnzbdSensor(SensorEntity):
def __init__(
self,
sabnzbd_api_data,
coordinator: SabnzbdUpdateCoordinator,
description: SabnzbdSensorEntityDescription,
entry_id,
) -> None:
"""Initialize the sensor."""
super().__init__(coordinator)
self._attr_unique_id = f"{entry_id}_{description.key}"
self.entity_description = description
self._sabnzbd_api = sabnzbd_api_data
self._attr_device_info = DeviceInfo(
entry_type=DeviceEntryType.SERVICE,
identifiers={(DOMAIN, entry_id)},
name=DEFAULT_NAME,
)
async def async_added_to_hass(self) -> None:
"""Call when entity about to be added to hass."""
self.async_on_remove(
async_dispatcher_connect(
self.hass, SIGNAL_SABNZBD_UPDATED, self.update_state
)
)
def update_state(self, args):
"""Get the latest data and updates the states."""
self._attr_native_value = self._sabnzbd_api.get_queue_field(
self.entity_description.key
)
if self._attr_native_value is not None:
if self.entity_description.key == SPEED_KEY:
self._attr_native_value = round(
float(self._attr_native_value) / 1024, 1
)
elif "size" in self.entity_description.key:
self._attr_native_value = round(float(self._attr_native_value), 2)
self.schedule_update_ha_state()
@property
def native_value(self) -> StateType:
"""Return latest sensor data."""
return self.coordinator.data.get(self.entity_description.key)