From 538a03ee0c97bc6a1da0749660e0c75e016e2150 Mon Sep 17 00:00:00 2001 From: Franck Nijhof Date: Fri, 28 May 2021 11:10:01 +0200 Subject: [PATCH] Clean up Speedtest.net Sensors (#51124) --- .../components/speedtestdotnet/sensor.py | 95 +++++++++---------- 1 file changed, 43 insertions(+), 52 deletions(-) diff --git a/homeassistant/components/speedtestdotnet/sensor.py b/homeassistant/components/speedtestdotnet/sensor.py index c49a5691cec..e28aa0b2527 100644 --- a/homeassistant/components/speedtestdotnet/sensor.py +++ b/homeassistant/components/speedtestdotnet/sensor.py @@ -1,7 +1,14 @@ """Support for Speedtest.net internet speed testing sensor.""" +from __future__ import annotations + +from typing import Any + from homeassistant.components.sensor import SensorEntity +from homeassistant.components.speedtestdotnet import SpeedTestDataCoordinator +from homeassistant.config_entries import ConfigEntry from homeassistant.const import ATTR_ATTRIBUTION -from homeassistant.core import callback +from homeassistant.core import HomeAssistant, callback +from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.restore_state import RestoreEntity from homeassistant.helpers.update_coordinator import CoordinatorEntity @@ -19,82 +26,64 @@ from .const import ( ) -async def async_setup_entry(hass, config_entry, async_add_entities): +async def async_setup_entry( + hass: HomeAssistant, + config_entry: ConfigEntry, + async_add_entities: AddEntitiesCallback, +) -> None: """Set up the Speedtestdotnet sensors.""" - speedtest_coordinator = hass.data[DOMAIN] - - entities = [] - for sensor_type in SENSOR_TYPES: - entities.append(SpeedtestSensor(speedtest_coordinator, sensor_type)) - - async_add_entities(entities) + async_add_entities( + SpeedtestSensor(speedtest_coordinator, sensor_type) + for sensor_type in SENSOR_TYPES + ) class SpeedtestSensor(CoordinatorEntity, RestoreEntity, SensorEntity): """Implementation of a speedtest.net sensor.""" - def __init__(self, coordinator, sensor_type): + coordinator: SpeedTestDataCoordinator + + _attr_icon = ICON + + def __init__(self, coordinator: SpeedTestDataCoordinator, sensor_type: str) -> None: """Initialize the sensor.""" super().__init__(coordinator) - self._name = SENSOR_TYPES[sensor_type][0] self.type = sensor_type - self._unit_of_measurement = SENSOR_TYPES[self.type][1] - self._state = None + + 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 @property - def name(self): - """Return the name of the sensor.""" - return f"{DEFAULT_NAME} {self._name}" - - @property - def unique_id(self): - """Return sensor unique_id.""" - return self.type - - @property - def state(self): - """Return the state of the device.""" - return self._state - - @property - def 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 - def extra_state_attributes(self): + def extra_state_attributes(self) -> dict[str, Any] | None: """Return the state attributes.""" if not self.coordinator.data: return None + attributes = { ATTR_ATTRIBUTION: ATTRIBUTION, ATTR_SERVER_NAME: self.coordinator.data["server"]["name"], ATTR_SERVER_COUNTRY: self.coordinator.data["server"]["country"], ATTR_SERVER_ID: self.coordinator.data["server"]["id"], } + if self.type == "download": attributes[ATTR_BYTES_RECEIVED] = self.coordinator.data["bytes_received"] - - if self.type == "upload": + elif self.type == "upload": attributes[ATTR_BYTES_SENT] = self.coordinator.data["bytes_sent"] return attributes - async def async_added_to_hass(self): + async def async_added_to_hass(self) -> None: """Handle entity which will be added.""" await super().async_added_to_hass() state = await self.async_get_last_state() if state: - self._state = state.state + self._attr_state = state.state @callback - def update(): + def update() -> None: """Update state.""" self._update_state() self.async_write_ha_state() @@ -102,12 +91,14 @@ class SpeedtestSensor(CoordinatorEntity, RestoreEntity, SensorEntity): self.async_on_remove(self.coordinator.async_add_listener(update)) self._update_state() - def _update_state(self): + def _update_state(self) -> None: """Update sensors state.""" - if self.coordinator.data: - if self.type == "ping": - self._state = self.coordinator.data["ping"] - elif self.type == "download": - self._state = round(self.coordinator.data["download"] / 10 ** 6, 2) - elif self.type == "upload": - self._state = round(self.coordinator.data["upload"] / 10 ** 6, 2) + if not self.coordinator.data: + return + + if self.type == "ping": + self._attr_state = self.coordinator.data["ping"] + elif self.type == "download": + self._attr_state = round(self.coordinator.data["download"] / 10 ** 6, 2) + elif self.type == "upload": + self._attr_state = round(self.coordinator.data["upload"] / 10 ** 6, 2)