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

View file

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

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."""
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)