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

.coveragerc

@@ -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/*

CODEOWNERS

@@ -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
+/homeassistant/components/sabnzbd/ @shaiu @jpbede
-/tests/components/sabnzbd/ @shaiu
+/tests/components/sabnzbd/ @shaiu @jpbede
 /homeassistant/components/saj/ @fredericvl
 /homeassistant/components/samsungtv/ @chemelli74 @epenet
 /tests/components/samsungtv/ @chemelli74 @epenet

homeassistant/components/sabnzbd/__init__.py

@@ -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:
+    async def async_pause_queue(
-        await api.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:
+    async def async_resume_queue(
-        await api.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)

homeassistant/components/sabnzbd/const.py

@@ -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"

homeassistant/components/sabnzbd/coordinator.py

@@ -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

homeassistant/components/sabnzbd/manifest.json

@@ -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",

homeassistant/components/sabnzbd/sensor.py

@@ -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 . import DOMAIN, SabnzbdUpdateCoordinator
-from .const import DEFAULT_NAME, KEY_API_DATA
+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
-
+    coordinator: SabnzbdUpdateCoordinator = hass.data[DOMAIN][entry_id]
-    sab_api_data = hass.data[DOMAIN][entry_id][KEY_API_DATA]
 
     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:
+    @property
-        """Call when entity about to be added to hass."""
+    def native_value(self) -> StateType:
-        self.async_on_remove(
+        """Return latest sensor data."""
-            async_dispatcher_connect(
+        return self.coordinator.data.get(self.entity_description.key)
-                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()