Use EntityDescription - iperf3 (#55681)

This commit is contained in:
Marc Mueller 2021-09-03 22:35:12 +02:00 committed by GitHub
parent 76ce0f6ea7
commit 1e4233fe20
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 39 additions and 46 deletions

View File

@ -1,11 +1,16 @@
"""Support for Iperf3 network measurement tool.""" """Support for Iperf3 network measurement tool."""
from __future__ import annotations
from datetime import timedelta from datetime import timedelta
import logging import logging
import iperf3 import iperf3
import voluptuous as vol import voluptuous as vol
from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN from homeassistant.components.sensor import (
DOMAIN as SENSOR_DOMAIN,
SensorEntityDescription,
)
from homeassistant.const import ( from homeassistant.const import (
CONF_HOST, CONF_HOST,
CONF_HOSTS, CONF_HOSTS,
@ -40,10 +45,19 @@ ATTR_UPLOAD = "upload"
ATTR_VERSION = "Version" ATTR_VERSION = "Version"
ATTR_HOST = "host" ATTR_HOST = "host"
SENSOR_TYPES = { SENSOR_TYPES: tuple[SensorEntityDescription, ...] = (
ATTR_DOWNLOAD: [ATTR_DOWNLOAD.capitalize(), DATA_RATE_MEGABITS_PER_SECOND], SensorEntityDescription(
ATTR_UPLOAD: [ATTR_UPLOAD.capitalize(), DATA_RATE_MEGABITS_PER_SECOND], key=ATTR_DOWNLOAD,
} name=ATTR_DOWNLOAD.capitalize(),
native_unit_of_measurement=DATA_RATE_MEGABITS_PER_SECOND,
),
SensorEntityDescription(
key=ATTR_UPLOAD,
name=ATTR_UPLOAD.capitalize(),
native_unit_of_measurement=DATA_RATE_MEGABITS_PER_SECOND,
),
)
SENSOR_KEYS: list[str] = [desc.key for desc in SENSOR_TYPES]
PROTOCOLS = ["tcp", "udp"] PROTOCOLS = ["tcp", "udp"]
@ -62,9 +76,9 @@ CONFIG_SCHEMA = vol.Schema(
DOMAIN: vol.Schema( DOMAIN: vol.Schema(
{ {
vol.Required(CONF_HOSTS): vol.All(cv.ensure_list, [HOST_CONFIG_SCHEMA]), vol.Required(CONF_HOSTS): vol.All(cv.ensure_list, [HOST_CONFIG_SCHEMA]),
vol.Optional( vol.Optional(CONF_MONITORED_CONDITIONS, default=SENSOR_KEYS): vol.All(
CONF_MONITORED_CONDITIONS, default=list(SENSOR_TYPES) cv.ensure_list, [vol.In(SENSOR_KEYS)]
): vol.All(cv.ensure_list, [vol.In(list(SENSOR_TYPES))]), ),
vol.Optional(CONF_SCAN_INTERVAL, default=DEFAULT_INTERVAL): vol.All( vol.Optional(CONF_SCAN_INTERVAL, default=DEFAULT_INTERVAL): vol.All(
cv.time_period, cv.positive_timedelta cv.time_period, cv.positive_timedelta
), ),

View File

@ -1,5 +1,5 @@
"""Support for Iperf3 sensors.""" """Support for Iperf3 sensors."""
from homeassistant.components.sensor import SensorEntity from homeassistant.components.sensor import SensorEntity, SensorEntityDescription
from homeassistant.const import ATTR_ATTRIBUTION from homeassistant.const import ATTR_ATTRIBUTION
from homeassistant.core import callback from homeassistant.core import callback
from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.dispatcher import async_dispatcher_connect
@ -18,42 +18,26 @@ ATTR_REMOTE_PORT = "Remote Port"
async def async_setup_platform(hass, config, async_add_entities, discovery_info): async def async_setup_platform(hass, config, async_add_entities, discovery_info):
"""Set up the Iperf3 sensor.""" """Set up the Iperf3 sensor."""
sensors = [] entities = [
for iperf3_host in hass.data[IPERF3_DOMAIN].values(): Iperf3Sensor(iperf3_host, description)
sensors.extend([Iperf3Sensor(iperf3_host, sensor) for sensor in discovery_info]) for iperf3_host in hass.data[IPERF3_DOMAIN].values()
async_add_entities(sensors, True) for description in SENSOR_TYPES
if description.key in discovery_info
]
async_add_entities(entities, True)
class Iperf3Sensor(RestoreEntity, SensorEntity): class Iperf3Sensor(RestoreEntity, SensorEntity):
"""A Iperf3 sensor implementation.""" """A Iperf3 sensor implementation."""
def __init__(self, iperf3_data, sensor_type): _attr_icon = ICON
_attr_should_poll = False
def __init__(self, iperf3_data, description: SensorEntityDescription):
"""Initialize the sensor.""" """Initialize the sensor."""
self._name = f"{SENSOR_TYPES[sensor_type][0]} {iperf3_data.host}" self.entity_description = description
self._state = None
self._sensor_type = sensor_type
self._unit_of_measurement = SENSOR_TYPES[sensor_type][1]
self._iperf3_data = iperf3_data self._iperf3_data = iperf3_data
self._attr_name = f"{description.name} {iperf3_data.host}"
@property
def name(self):
"""Return the name of the sensor."""
return self._name
@property
def native_value(self):
"""Return the state of the device."""
return self._state
@property
def native_unit_of_measurement(self):
"""Return the unit of measurement of this entity, if any."""
return self._unit_of_measurement
@property
def icon(self):
"""Return icon."""
return ICON
@property @property
def extra_state_attributes(self): def extra_state_attributes(self):
@ -66,11 +50,6 @@ class Iperf3Sensor(RestoreEntity, SensorEntity):
ATTR_VERSION: self._iperf3_data.data[ATTR_VERSION], ATTR_VERSION: self._iperf3_data.data[ATTR_VERSION],
} }
@property
def should_poll(self):
"""Return the polling requirement for this sensor."""
return False
async def async_added_to_hass(self): async def async_added_to_hass(self):
"""Handle entity which will be added.""" """Handle entity which will be added."""
await super().async_added_to_hass() await super().async_added_to_hass()
@ -84,13 +63,13 @@ class Iperf3Sensor(RestoreEntity, SensorEntity):
state = await self.async_get_last_state() state = await self.async_get_last_state()
if not state: if not state:
return return
self._state = state.state self._attr_native_value = state.state
def update(self): def update(self):
"""Get the latest data and update the states.""" """Get the latest data and update the states."""
data = self._iperf3_data.data.get(self._sensor_type) data = self._iperf3_data.data.get(self.entity_description.key)
if data is not None: if data is not None:
self._state = round(data, 2) self._attr_native_value = round(data, 2)
@callback @callback
def _schedule_immediate_update(self, host): def _schedule_immediate_update(self, host):