diff --git a/homeassistant/components/iperf3/__init__.py b/homeassistant/components/iperf3/__init__.py index bd5aeac099a..04dabc013e7 100644 --- a/homeassistant/components/iperf3/__init__.py +++ b/homeassistant/components/iperf3/__init__.py @@ -1,11 +1,16 @@ """Support for Iperf3 network measurement tool.""" +from __future__ import annotations + from datetime import timedelta import logging import iperf3 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 ( CONF_HOST, CONF_HOSTS, @@ -40,10 +45,19 @@ ATTR_UPLOAD = "upload" ATTR_VERSION = "Version" ATTR_HOST = "host" -SENSOR_TYPES = { - ATTR_DOWNLOAD: [ATTR_DOWNLOAD.capitalize(), DATA_RATE_MEGABITS_PER_SECOND], - ATTR_UPLOAD: [ATTR_UPLOAD.capitalize(), DATA_RATE_MEGABITS_PER_SECOND], -} +SENSOR_TYPES: tuple[SensorEntityDescription, ...] = ( + SensorEntityDescription( + 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"] @@ -62,9 +76,9 @@ CONFIG_SCHEMA = vol.Schema( DOMAIN: vol.Schema( { vol.Required(CONF_HOSTS): vol.All(cv.ensure_list, [HOST_CONFIG_SCHEMA]), - vol.Optional( - CONF_MONITORED_CONDITIONS, default=list(SENSOR_TYPES) - ): vol.All(cv.ensure_list, [vol.In(list(SENSOR_TYPES))]), + vol.Optional(CONF_MONITORED_CONDITIONS, default=SENSOR_KEYS): vol.All( + cv.ensure_list, [vol.In(SENSOR_KEYS)] + ), vol.Optional(CONF_SCAN_INTERVAL, default=DEFAULT_INTERVAL): vol.All( cv.time_period, cv.positive_timedelta ), diff --git a/homeassistant/components/iperf3/sensor.py b/homeassistant/components/iperf3/sensor.py index 07b9cc069e4..dfc4abba707 100644 --- a/homeassistant/components/iperf3/sensor.py +++ b/homeassistant/components/iperf3/sensor.py @@ -1,5 +1,5 @@ """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.core import callback 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): """Set up the Iperf3 sensor.""" - sensors = [] - for iperf3_host in hass.data[IPERF3_DOMAIN].values(): - sensors.extend([Iperf3Sensor(iperf3_host, sensor) for sensor in discovery_info]) - async_add_entities(sensors, True) + entities = [ + Iperf3Sensor(iperf3_host, description) + for iperf3_host in hass.data[IPERF3_DOMAIN].values() + for description in SENSOR_TYPES + if description.key in discovery_info + ] + async_add_entities(entities, True) class Iperf3Sensor(RestoreEntity, SensorEntity): """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.""" - self._name = f"{SENSOR_TYPES[sensor_type][0]} {iperf3_data.host}" - self._state = None - self._sensor_type = sensor_type - self._unit_of_measurement = SENSOR_TYPES[sensor_type][1] + self.entity_description = description self._iperf3_data = iperf3_data - - @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 + self._attr_name = f"{description.name} {iperf3_data.host}" @property def extra_state_attributes(self): @@ -66,11 +50,6 @@ class Iperf3Sensor(RestoreEntity, SensorEntity): 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): """Handle entity which will be added.""" await super().async_added_to_hass() @@ -84,13 +63,13 @@ class Iperf3Sensor(RestoreEntity, SensorEntity): state = await self.async_get_last_state() if not state: return - self._state = state.state + self._attr_native_value = state.state def update(self): """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: - self._state = round(data, 2) + self._attr_native_value = round(data, 2) @callback def _schedule_immediate_update(self, host):