Add Update coordinator to QBittorrent (#98896)

.coveragerc

@@ -989,6 +989,7 @@ omit =
     homeassistant/components/pushsafer/notify.py
     homeassistant/components/pyload/sensor.py
     homeassistant/components/qbittorrent/__init__.py
+    homeassistant/components/qbittorrent/coordinator.py
     homeassistant/components/qbittorrent/sensor.py
     homeassistant/components/qnap/__init__.py
     homeassistant/components/qnap/coordinator.py

homeassistant/components/qbittorrent/__init__.py

@@ -16,6 +16,7 @@ from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
 
 from .const import DOMAIN
+from .coordinator import QBittorrentDataCoordinator
 from .helpers import setup_client
 
 PLATFORMS = [Platform.SENSOR]
@@ -27,7 +28,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up qBittorrent from a config entry."""
     hass.data.setdefault(DOMAIN, {})
     try:
-        hass.data[DOMAIN][entry.entry_id] = await hass.async_add_executor_job(
+        client = await hass.async_add_executor_job(
             setup_client,
             entry.data[CONF_URL],
             entry.data[CONF_USERNAME],
@@ -38,7 +39,11 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         raise ConfigEntryNotReady("Invalid credentials") from err
     except RequestException as err:
         raise ConfigEntryNotReady("Failed to connect") from err
+    coordinator = QBittorrentDataCoordinator(hass, client)
 
+    await coordinator.async_config_entry_first_refresh()
+
+    hass.data[DOMAIN][entry.entry_id] = coordinator
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
     return True
 

homeassistant/components/qbittorrent/coordinator.py

@@ -0,0 +1,38 @@
+"""The QBittorrent coordinator."""
+from __future__ import annotations
+
+from datetime import timedelta
+import logging
+from typing import Any
+
+from qbittorrent import Client
+from qbittorrent.client import LoginRequired
+
+from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import ConfigEntryError
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
+
+from .const import DOMAIN
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class QBittorrentDataCoordinator(DataUpdateCoordinator[dict[str, Any]]):
+    """QBittorrent update coordinator."""
+
+    def __init__(self, hass: HomeAssistant, client: Client) -> None:
+        """Initialize coordinator."""
+        self.client = client
+
+        super().__init__(
+            hass,
+            _LOGGER,
+            name=DOMAIN,
+            update_interval=timedelta(seconds=30),
+        )
+
+    async def _async_update_data(self) -> dict[str, Any]:
+        try:
+            return await self.hass.async_add_executor_job(self.client.sync_main_data)
+        except LoginRequired as exc:
+            raise ConfigEntryError("Invalid authentication") from exc

homeassistant/components/qbittorrent/sensor.py

@@ -1,10 +1,10 @@
 """Support for monitoring the qBittorrent API."""
 from __future__ import annotations
 
+from collections.abc import Callable
+from dataclasses import dataclass
 import logging
-
-from qbittorrent.client import Client, LoginRequired
-from requests.exceptions import RequestException
+from typing import Any
 
 from homeassistant.components.sensor import (
     SensorDeviceClass,
@@ -16,8 +16,11 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import STATE_IDLE, UnitOfDataRate
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.helpers.typing import StateType
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from .const import DOMAIN
+from .coordinator import QBittorrentDataCoordinator
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -25,26 +28,61 @@ SENSOR_TYPE_CURRENT_STATUS = "current_status"
 SENSOR_TYPE_DOWNLOAD_SPEED = "download_speed"
 SENSOR_TYPE_UPLOAD_SPEED = "upload_speed"
 
-SENSOR_TYPES: tuple[SensorEntityDescription, ...] = (
-    SensorEntityDescription(
+
+@dataclass
+class QBittorrentMixin:
+    """Mixin for required keys."""
+
+    value_fn: Callable[[dict[str, Any]], StateType]
+
+
+@dataclass
+class QBittorrentSensorEntityDescription(SensorEntityDescription, QBittorrentMixin):
+    """Describes QBittorrent sensor entity."""
+
+
+def _get_qbittorrent_state(data: dict[str, Any]) -> str:
+    download = data["server_state"]["dl_info_speed"]
+    upload = data["server_state"]["up_info_speed"]
+
+    if upload > 0 and download > 0:
+        return "up_down"
+    if upload > 0 and download == 0:
+        return "seeding"
+    if upload == 0 and download > 0:
+        return "downloading"
+    return STATE_IDLE
+
+
+def format_speed(speed):
+    """Return a bytes/s measurement as a human readable string."""
+    kb_spd = float(speed) / 1024
+    return round(kb_spd, 2 if kb_spd < 0.1 else 1)
+
+
+SENSOR_TYPES: tuple[QBittorrentSensorEntityDescription, ...] = (
+    QBittorrentSensorEntityDescription(
         key=SENSOR_TYPE_CURRENT_STATUS,
         name="Status",
+        value_fn=_get_qbittorrent_state,
     ),
-    SensorEntityDescription(
+    QBittorrentSensorEntityDescription(
         key=SENSOR_TYPE_DOWNLOAD_SPEED,
         name="Down Speed",
         icon="mdi:cloud-download",
         device_class=SensorDeviceClass.DATA_RATE,
         native_unit_of_measurement=UnitOfDataRate.KIBIBYTES_PER_SECOND,
         state_class=SensorStateClass.MEASUREMENT,
+        value_fn=lambda data: format_speed(data["server_state"]["dl_info_speed"]),
     ),
-    SensorEntityDescription(
+    QBittorrentSensorEntityDescription(
         key=SENSOR_TYPE_UPLOAD_SPEED,
         name="Up Speed",
         icon="mdi:cloud-upload",
         device_class=SensorDeviceClass.DATA_RATE,
         native_unit_of_measurement=UnitOfDataRate.KIBIBYTES_PER_SECOND,
         state_class=SensorStateClass.MEASUREMENT,
+        value_fn=lambda data: format_speed(data["server_state"]["up_info_speed"]),
     ),
 )
 
@@ -55,68 +93,33 @@ async def async_setup_entry(
     async_add_entites: AddEntitiesCallback,
 ) -> None:
     """Set up qBittorrent sensor entries."""
-    client: Client = hass.data[DOMAIN][config_entry.entry_id]
+    coordinator: QBittorrentDataCoordinator = hass.data[DOMAIN][config_entry.entry_id]
     entities = [
-        QBittorrentSensor(description, client, config_entry)
+        QBittorrentSensor(description, coordinator, config_entry)
         for description in SENSOR_TYPES
     ]
-    async_add_entites(entities, True)
+    async_add_entites(entities)
 
 
-def format_speed(speed):
-    """Return a bytes/s measurement as a human readable string."""
-    kb_spd = float(speed) / 1024
-    return round(kb_spd, 2 if kb_spd < 0.1 else 1)
+class QBittorrentSensor(CoordinatorEntity[QBittorrentDataCoordinator], SensorEntity):
+    """Representation of a qBittorrent sensor."""
 
-
-class QBittorrentSensor(SensorEntity):
-    """Representation of an qBittorrent sensor."""
+    entity_description: QBittorrentSensorEntityDescription
 
     def __init__(
         self,
-        description: SensorEntityDescription,
-        qbittorrent_client: Client,
+        description: QBittorrentSensorEntityDescription,
+        coordinator: QBittorrentDataCoordinator,
         config_entry: ConfigEntry,
     ) -> None:
         """Initialize the qBittorrent sensor."""
+        super().__init__(coordinator)
         self.entity_description = description
-        self.client = qbittorrent_client
-
         self._attr_unique_id = f"{config_entry.entry_id}-{description.key}"
         self._attr_name = f"{config_entry.title} {description.name}"
         self._attr_available = False
 
-    def update(self) -> None:
-        """Get the latest data from qBittorrent and updates the state."""
-        try:
-            data = self.client.sync_main_data()
-            self._attr_available = True
-        except RequestException:
-            _LOGGER.error("Connection lost")
-            self._attr_available = False
-            return
-        except LoginRequired:
-            _LOGGER.error("Invalid authentication")
-            return
-
-        if data is None:
-            return
-
-        download = data["server_state"]["dl_info_speed"]
-        upload = data["server_state"]["up_info_speed"]
-
-        sensor_type = self.entity_description.key
-        if sensor_type == SENSOR_TYPE_CURRENT_STATUS:
-            if upload > 0 and download > 0:
-                self._attr_native_value = "up_down"
-            elif upload > 0 and download == 0:
-                self._attr_native_value = "seeding"
-            elif upload == 0 and download > 0:
-                self._attr_native_value = "downloading"
-            else:
-                self._attr_native_value = STATE_IDLE
-
-        elif sensor_type == SENSOR_TYPE_DOWNLOAD_SPEED:
-            self._attr_native_value = format_speed(download)
-        elif sensor_type == SENSOR_TYPE_UPLOAD_SPEED:
-            self._attr_native_value = format_speed(upload)
+    @property
+    def native_value(self) -> StateType:
+        """Return value of sensor."""
+        return self.entity_description.value_fn(self.coordinator.data)