Add Update coordinator to QBittorrent (#98896)

This commit is contained in:
Joost Lekkerkerker 2023-10-25 15:51:52 +02:00 committed by GitHub
parent 8d034a85fe
commit cd8e3a81db
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 103 additions and 56 deletions

View file

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

View file

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

View file

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

View file

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