From bffa9f960dac0130f26d98fc16fce67341eedf30 Mon Sep 17 00:00:00 2001 From: Michael <35783820+mib1185@users.noreply.github.com> Date: Mon, 2 Aug 2021 17:00:25 +0200 Subject: [PATCH] Add state class measurement to all suitable sensors on Speedtest.net (#53693) * Add state class measurement * use tuple instead of list --- .../components/speedtestdotnet/const.py | 31 ++++++++++++++++--- .../components/speedtestdotnet/sensor.py | 29 +++++++++-------- .../components/speedtestdotnet/test_sensor.py | 8 ++--- 3 files changed, 45 insertions(+), 23 deletions(-) diff --git a/homeassistant/components/speedtestdotnet/const.py b/homeassistant/components/speedtestdotnet/const.py index 04f3ea0cc55..897ffa126fa 100644 --- a/homeassistant/components/speedtestdotnet/const.py +++ b/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" diff --git a/homeassistant/components/speedtestdotnet/sensor.py b/homeassistant/components/speedtestdotnet/sensor.py index 8dcc5bc3459..1c6c80a6af1 100644 --- a/homeassistant/components/speedtestdotnet/sensor.py +++ b/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) diff --git a/tests/components/speedtestdotnet/test_sensor.py b/tests/components/speedtestdotnet/test_sensor.py index 11db05d2994..d0378731c28 100644 --- a/tests/components/speedtestdotnet/test_sensor.py +++ b/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]