Use EntityDescription - iperf3 (#55681)

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
                 ),

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):