Use DataUpdate coordinator for Transmission (#99209)

* Switch integration to DataUpdate Coordinator

* add coordinator to .coveragerc

* Migrate TransmissionData into DUC

* update coveragerc

* Applu suggestions

* remove CONFIG_SCHEMA
This commit is contained in:
Rami Mosleh 2023-10-12 21:58:22 +03:00 committed by GitHub
parent cc3d1a11bd
commit 536ad57bf4
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 340 additions and 465 deletions

View file

@ -4,21 +4,18 @@ from __future__ import annotations
from contextlib import suppress
from typing import Any
from transmission_rpc.session import SessionStats
from transmission_rpc.torrent import Torrent
from homeassistant.components.sensor import SensorDeviceClass, SensorEntity
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_NAME, STATE_IDLE, UnitOfDataRate
from homeassistant.core import HomeAssistant, callback
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 TransmissionClient
from .const import (
CONF_LIMIT,
CONF_ORDER,
DOMAIN,
STATE_ATTR_TORRENT_INFO,
STATE_DOWNLOADING,
@ -26,6 +23,7 @@ from .const import (
STATE_UP_DOWN,
SUPPORTED_ORDER_MODES,
)
from .coordinator import TransmissionDataUpdateCoordinator
async def async_setup_entry(
@ -35,54 +33,56 @@ async def async_setup_entry(
) -> None:
"""Set up the Transmission sensors."""
tm_client: TransmissionClient = hass.data[DOMAIN][config_entry.entry_id]
coordinator: TransmissionDataUpdateCoordinator = hass.data[DOMAIN][
config_entry.entry_id
]
name: str = config_entry.data[CONF_NAME]
dev = [
TransmissionSpeedSensor(
tm_client,
coordinator,
name,
"download_speed",
"download",
),
TransmissionSpeedSensor(
tm_client,
coordinator,
name,
"upload_speed",
"upload",
),
TransmissionStatusSensor(
tm_client,
coordinator,
name,
"transmission_status",
"status",
),
TransmissionTorrentsSensor(
tm_client,
coordinator,
name,
"active_torrents",
"active_torrents",
),
TransmissionTorrentsSensor(
tm_client,
coordinator,
name,
"paused_torrents",
"paused_torrents",
),
TransmissionTorrentsSensor(
tm_client,
coordinator,
name,
"total_torrents",
"total_torrents",
),
TransmissionTorrentsSensor(
tm_client,
coordinator,
name,
"completed_torrents",
"completed_torrents",
),
TransmissionTorrentsSensor(
tm_client,
coordinator,
name,
"started_torrents",
"started_torrents",
@ -92,7 +92,7 @@ async def async_setup_entry(
async_add_entities(dev, True)
class TransmissionSensor(SensorEntity):
class TransmissionSensor(CoordinatorEntity[SessionStats], SensorEntity):
"""A base class for all Transmission sensors."""
_attr_has_entity_name = True
@ -100,48 +100,23 @@ class TransmissionSensor(SensorEntity):
def __init__(
self,
tm_client: TransmissionClient,
coordinator: TransmissionDataUpdateCoordinator,
client_name: str,
sensor_translation_key: str,
key: str,
) -> None:
"""Initialize the sensor."""
self._tm_client = tm_client
super().__init__(coordinator)
self._attr_translation_key = sensor_translation_key
self._key = key
self._state: StateType = None
self._attr_unique_id = f"{tm_client.config_entry.entry_id}-{key}"
self._attr_unique_id = f"{coordinator.config_entry.entry_id}-{key}"
self._attr_device_info = DeviceInfo(
entry_type=DeviceEntryType.SERVICE,
identifiers={(DOMAIN, tm_client.config_entry.entry_id)},
identifiers={(DOMAIN, coordinator.config_entry.entry_id)},
manufacturer="Transmission",
name=client_name,
)
@property
def native_value(self) -> StateType:
"""Return the state of the sensor."""
return self._state
@property
def available(self) -> bool:
"""Could the device be accessed during the last update call."""
return self._tm_client.api.available
async def async_added_to_hass(self) -> None:
"""Handle entity which will be added."""
@callback
def update():
"""Update the state."""
self.async_schedule_update_ha_state(True)
self.async_on_remove(
async_dispatcher_connect(
self.hass, self._tm_client.api.signal_update, update
)
)
class TransmissionSpeedSensor(TransmissionSensor):
"""Representation of a Transmission speed sensor."""
@ -151,15 +126,15 @@ class TransmissionSpeedSensor(TransmissionSensor):
_attr_suggested_display_precision = 2
_attr_suggested_unit_of_measurement = UnitOfDataRate.MEGABYTES_PER_SECOND
def update(self) -> None:
"""Get the latest data from Transmission and updates the state."""
if data := self._tm_client.api.data:
b_spd = (
float(data.download_speed)
if self._key == "download"
else float(data.upload_speed)
)
self._state = b_spd
@property
def native_value(self) -> float:
"""Return the speed of the sensor."""
data = self.coordinator.data
return (
float(data.download_speed)
if self._key == "download"
else float(data.upload_speed)
)
class TransmissionStatusSensor(TransmissionSensor):
@ -168,21 +143,18 @@ class TransmissionStatusSensor(TransmissionSensor):
_attr_device_class = SensorDeviceClass.ENUM
_attr_options = [STATE_IDLE, STATE_UP_DOWN, STATE_SEEDING, STATE_DOWNLOADING]
def update(self) -> None:
"""Get the latest data from Transmission and updates the state."""
if data := self._tm_client.api.data:
upload = data.upload_speed
download = data.download_speed
if upload > 0 and download > 0:
self._state = STATE_UP_DOWN
elif upload > 0 and download == 0:
self._state = STATE_SEEDING
elif upload == 0 and download > 0:
self._state = STATE_DOWNLOADING
else:
self._state = STATE_IDLE
else:
self._state = None
@property
def native_value(self) -> str:
"""Return the value of the status sensor."""
upload = self.coordinator.data.upload_speed
download = self.coordinator.data.download_speed
if upload > 0 and download > 0:
return STATE_UP_DOWN
if upload > 0 and download == 0:
return STATE_SEEDING
if upload == 0 and download > 0:
return STATE_DOWNLOADING
return STATE_IDLE
class TransmissionTorrentsSensor(TransmissionSensor):
@ -208,21 +180,22 @@ class TransmissionTorrentsSensor(TransmissionSensor):
def extra_state_attributes(self) -> dict[str, Any]:
"""Return the state attributes, if any."""
info = _torrents_info(
torrents=self._tm_client.api.torrents,
order=self._tm_client.config_entry.options[CONF_ORDER],
limit=self._tm_client.config_entry.options[CONF_LIMIT],
torrents=self.coordinator.torrents,
order=self.coordinator.order,
limit=self.coordinator.limit,
statuses=self.MODES[self._key],
)
return {
STATE_ATTR_TORRENT_INFO: info,
}
def update(self) -> None:
"""Get the latest data from Transmission and updates the state."""
@property
def native_value(self) -> int:
"""Return the count of the sensor."""
torrents = _filter_torrents(
self._tm_client.api.torrents, statuses=self.MODES[self._key]
self.coordinator.torrents, statuses=self.MODES[self._key]
)
self._state = len(torrents)
return len(torrents)
def _filter_torrents(