From cd8e3a81dbf8827de675f4e24bc402dfef56f0d5 Mon Sep 17 00:00:00 2001 From: Joost Lekkerkerker Date: Wed, 25 Oct 2023 15:51:52 +0200 Subject: [PATCH] Add Update coordinator to QBittorrent (#98896) --- .coveragerc | 1 + .../components/qbittorrent/__init__.py | 7 +- .../components/qbittorrent/coordinator.py | 38 ++++++ .../components/qbittorrent/sensor.py | 113 +++++++++--------- 4 files changed, 103 insertions(+), 56 deletions(-) create mode 100644 homeassistant/components/qbittorrent/coordinator.py diff --git a/.coveragerc b/.coveragerc index 86b92c07a3d..5ef7ece3bd8 100644 --- a/.coveragerc +++ b/.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 diff --git a/homeassistant/components/qbittorrent/__init__.py b/homeassistant/components/qbittorrent/__init__.py index 53e8d4b9660..fd9577f5c73 100644 --- a/homeassistant/components/qbittorrent/__init__.py +++ b/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 diff --git a/homeassistant/components/qbittorrent/coordinator.py b/homeassistant/components/qbittorrent/coordinator.py new file mode 100644 index 00000000000..8363a764d0a --- /dev/null +++ b/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 diff --git a/homeassistant/components/qbittorrent/sensor.py b/homeassistant/components/qbittorrent/sensor.py index 5cca77ecc34..e2feee1e60c 100644 --- a/homeassistant/components/qbittorrent/sensor.py +++ b/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)