Add state class measurement to all suitable sensors on Speedtest.net (#53693)

* Add state class measurement

* use tuple instead of list

homeassistant/components/speedtestdotnet/const.py

@@ -1,17 +1,38 @@
 """Consts used by Speedtest.net."""
+from __future__ import annotations
+
 from typing import Final
 
+from homeassistant.components.sensor import (
+    STATE_CLASS_MEASUREMENT,
+    SensorEntityDescription,
+)
 from homeassistant.const import DATA_RATE_MEGABITS_PER_SECOND, TIME_MILLISECONDS
 
 DOMAIN: Final = "speedtestdotnet"
 
 SPEED_TEST_SERVICE: Final = "speedtest"
 
-SENSOR_TYPES: Final = {
-    "ping": ["Ping", TIME_MILLISECONDS],
-    "download": ["Download", DATA_RATE_MEGABITS_PER_SECOND],
-    "upload": ["Upload", DATA_RATE_MEGABITS_PER_SECOND],
-}
+SENSOR_TYPES: Final[tuple[SensorEntityDescription, ...]] = (
+    SensorEntityDescription(
+        key="ping",
+        name="Ping",
+        unit_of_measurement=TIME_MILLISECONDS,
+        state_class=STATE_CLASS_MEASUREMENT,
+    ),
+    SensorEntityDescription(
+        key="download",
+        name="Download",
+        unit_of_measurement=DATA_RATE_MEGABITS_PER_SECOND,
+        state_class=STATE_CLASS_MEASUREMENT,
+    ),
+    SensorEntityDescription(
+        key="upload",
+        name="Upload",
+        unit_of_measurement=DATA_RATE_MEGABITS_PER_SECOND,
+        state_class=STATE_CLASS_MEASUREMENT,
+    ),
+)
 
 CONF_SERVER_NAME: Final = "server_name"
 CONF_SERVER_ID: Final = "server_id"

homeassistant/components/speedtestdotnet/sensor.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 
 from typing import Any
 
-from homeassistant.components.sensor import SensorEntity
+from homeassistant.components.sensor import SensorEntity, SensorEntityDescription
 from homeassistant.components.speedtestdotnet import SpeedTestDataCoordinator
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_ATTRIBUTION
@@ -34,8 +34,8 @@ async def async_setup_entry(
     """Set up the Speedtestdotnet sensors."""
     speedtest_coordinator = hass.data[DOMAIN]
     async_add_entities(
-        SpeedtestSensor(speedtest_coordinator, sensor_type)
-        for sensor_type in SENSOR_TYPES
+        SpeedtestSensor(speedtest_coordinator, description)
+        for description in SENSOR_TYPES
     )
 
 
@@ -46,14 +46,17 @@ class SpeedtestSensor(CoordinatorEntity, RestoreEntity, SensorEntity):
 
     _attr_icon = ICON
 
-    def __init__(self, coordinator: SpeedTestDataCoordinator, sensor_type: str) -> None:
+    def __init__(
+        self,
+        coordinator: SpeedTestDataCoordinator,
+        description: SensorEntityDescription,
+    ) -> None:
         """Initialize the sensor."""
         super().__init__(coordinator)
-        self.type = sensor_type
+        self.entity_description = description
 
-        self._attr_name = f"{DEFAULT_NAME} {SENSOR_TYPES[sensor_type][0]}"
-        self._attr_unit_of_measurement = SENSOR_TYPES[self.type][1]
-        self._attr_unique_id = sensor_type
+        self._attr_name = f"{DEFAULT_NAME} {description.name}"
+        self._attr_unique_id = description.key
         self._attrs = {ATTR_ATTRIBUTION: ATTRIBUTION}
 
     @property
@@ -68,11 +71,11 @@ class SpeedtestSensor(CoordinatorEntity, RestoreEntity, SensorEntity):
                 }
             )
 
-            if self.type == "download":
+            if self.entity_description.key == "download":
                 self._attrs[ATTR_BYTES_RECEIVED] = self.coordinator.data[
                     "bytes_received"
                 ]
-            elif self.type == "upload":
+            elif self.entity_description.key == "upload":
                 self._attrs[ATTR_BYTES_SENT] = self.coordinator.data["bytes_sent"]
 
         return self._attrs
@@ -96,9 +99,9 @@ class SpeedtestSensor(CoordinatorEntity, RestoreEntity, SensorEntity):
     def _update_state(self):
         """Update sensors state."""
         if self.coordinator.data:
-            if self.type == "ping":
+            if self.entity_description.key == "ping":
                 self._attr_state = self.coordinator.data["ping"]
-            elif self.type == "download":
+            elif self.entity_description.key == "download":
                 self._attr_state = round(self.coordinator.data["download"] / 10 ** 6, 2)
-            elif self.type == "upload":
+            elif self.entity_description.key == "upload":
                 self._attr_state = round(self.coordinator.data["upload"] / 10 ** 6, 2)

tests/components/speedtestdotnet/test_sensor.py

@@ -26,9 +26,7 @@ async def test_speedtestdotnet_sensors(
 
     assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 3
 
-    for sensor_type in SENSOR_TYPES:
-        sensor = hass.states.get(
-            f"sensor.{DEFAULT_NAME}_{SENSOR_TYPES[sensor_type][0]}"
-        )
+    for description in SENSOR_TYPES:
+        sensor = hass.states.get(f"sensor.{DEFAULT_NAME}_{description.name}")
         assert sensor
-        assert sensor.state == MOCK_STATES[sensor_type]
+        assert sensor.state == MOCK_STATES[description.key]