From 10a6c56fec9d3a960d315587df351190f831f724 Mon Sep 17 00:00:00 2001 From: Rami Mosleh Date: Sun, 18 Dec 2022 11:57:17 +0200 Subject: [PATCH] Move speedtestdotnet coordinator to separate file (#83979) --- .../components/speedtestdotnet/__init__.py | 133 ++++-------------- .../components/speedtestdotnet/config_flow.py | 2 +- .../components/speedtestdotnet/const.py | 44 ------ .../components/speedtestdotnet/coordinator.py | 79 +++++++++++ .../components/speedtestdotnet/sensor.py | 44 +++++- .../components/speedtestdotnet/test_sensor.py | 15 +- 6 files changed, 155 insertions(+), 162 deletions(-) create mode 100644 homeassistant/components/speedtestdotnet/coordinator.py diff --git a/homeassistant/components/speedtestdotnet/__init__.py b/homeassistant/components/speedtestdotnet/__init__.py index 8cf32def197..3f901a914c5 100644 --- a/homeassistant/components/speedtestdotnet/__init__.py +++ b/homeassistant/components/speedtestdotnet/__init__.py @@ -2,34 +2,37 @@ from __future__ import annotations from datetime import timedelta -import logging +from functools import partial import speedtest from homeassistant.config_entries import ConfigEntry -from homeassistant.const import CONF_SCAN_INTERVAL, EVENT_HOMEASSISTANT_STARTED -from homeassistant.core import CoreState, HomeAssistant -from homeassistant.exceptions import ConfigEntryNotReady -from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed - -from .const import ( - CONF_MANUAL, - CONF_SERVER_ID, - DEFAULT_SCAN_INTERVAL, - DEFAULT_SERVER, - DOMAIN, - PLATFORMS, +from homeassistant.const import ( + CONF_SCAN_INTERVAL, + EVENT_HOMEASSISTANT_STARTED, + Platform, ) +from homeassistant.core import CoreState, Event, HomeAssistant +from homeassistant.exceptions import ConfigEntryNotReady -_LOGGER = logging.getLogger(__name__) +from .const import CONF_MANUAL, DEFAULT_SCAN_INTERVAL, DOMAIN +from .coordinator import SpeedTestDataCoordinator + +PLATFORMS = [Platform.SENSOR] async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool: """Set up the Speedtest.net component.""" - coordinator = SpeedTestDataCoordinator(hass, config_entry) - await coordinator.async_setup() + try: + api = await hass.async_add_executor_job( + partial(speedtest.Speedtest, secure=True) + ) + coordinator = SpeedTestDataCoordinator(hass, config_entry, api) + await hass.async_add_executor_job(coordinator.update_servers) + except speedtest.SpeedtestException as err: + raise ConfigEntryNotReady from err - async def _enable_scheduled_speedtests(*_): + async def _enable_scheduled_speedtests(event: Event | None = None) -> None: """Activate the data update coordinator.""" coordinator.update_interval = timedelta( minutes=config_entry.options.get(CONF_SCAN_INTERVAL, DEFAULT_SCAN_INTERVAL) @@ -51,104 +54,28 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b await hass.config_entries.async_forward_entry_setups(config_entry, PLATFORMS) + config_entry.async_on_unload( + config_entry.add_update_listener(options_updated_listener) + ) + return True async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool: """Unload SpeedTest Entry from config_entry.""" - unload_ok = await hass.config_entries.async_unload_platforms( + if unload_ok := await hass.config_entries.async_unload_platforms( config_entry, PLATFORMS - ) - if unload_ok: + ): hass.data.pop(DOMAIN) return unload_ok -class SpeedTestDataCoordinator(DataUpdateCoordinator): - """Get the latest data from speedtest.net.""" - - def __init__(self, hass: HomeAssistant, config_entry: ConfigEntry) -> None: - """Initialize the data object.""" - self.hass = hass - self.config_entry: ConfigEntry = config_entry - self.api: speedtest.Speedtest | None = None - self.servers: dict[str, dict] = {DEFAULT_SERVER: {}} - super().__init__( - self.hass, - _LOGGER, - name=DOMAIN, - update_method=self.async_update, - ) - - def initialize(self) -> None: - """Initialize speedtest api.""" - self.api = speedtest.Speedtest(secure=True) - self.update_servers() - - def update_servers(self): - """Update list of test servers.""" - test_servers = self.api.get_servers() - test_servers_list = [] - for servers in test_servers.values(): - for server in servers: - test_servers_list.append(server) - for server in sorted( - test_servers_list, - key=lambda server: ( - server["country"], - server["name"], - server["sponsor"], - ), - ): - self.servers[ - f"{server['country']} - {server['sponsor']} - {server['name']}" - ] = server - - def update_data(self): - """Get the latest data from speedtest.net.""" - self.update_servers() - self.api.closest.clear() - if self.config_entry.options.get(CONF_SERVER_ID): - server_id = self.config_entry.options.get(CONF_SERVER_ID) - self.api.get_servers(servers=[server_id]) - - best_server = self.api.get_best_server() - _LOGGER.debug( - "Executing speedtest.net speed test with server_id: %s", - best_server["id"], - ) - self.api.download() - self.api.upload() - return self.api.results.dict() - - async def async_update(self) -> dict[str, str]: - """Update Speedtest data.""" - try: - return await self.hass.async_add_executor_job(self.update_data) - except speedtest.NoMatchedServers as err: - raise UpdateFailed("Selected server is not found.") from err - except speedtest.SpeedtestException as err: - raise UpdateFailed(err) from err - - async def async_setup(self) -> None: - """Set up SpeedTest.""" - try: - await self.hass.async_add_executor_job(self.initialize) - except speedtest.SpeedtestException as err: - raise ConfigEntryNotReady from err - - self.config_entry.async_on_unload( - self.config_entry.add_update_listener(options_updated_listener) - ) - - async def options_updated_listener(hass: HomeAssistant, entry: ConfigEntry) -> None: """Handle options update.""" + coordinator: SpeedTestDataCoordinator = hass.data[DOMAIN] if entry.options[CONF_MANUAL]: - hass.data[DOMAIN].update_interval = None + coordinator.update_interval = None return - hass.data[DOMAIN].update_interval = timedelta( - minutes=entry.options[CONF_SCAN_INTERVAL] - ) - await hass.data[DOMAIN].async_request_refresh() + coordinator.update_interval = timedelta(minutes=entry.options[CONF_SCAN_INTERVAL]) + await coordinator.async_request_refresh() diff --git a/homeassistant/components/speedtestdotnet/config_flow.py b/homeassistant/components/speedtestdotnet/config_flow.py index d82ac6bf728..e6a0ff7bcb2 100644 --- a/homeassistant/components/speedtestdotnet/config_flow.py +++ b/homeassistant/components/speedtestdotnet/config_flow.py @@ -30,7 +30,7 @@ class SpeedTestFlowHandler(config_entries.ConfigFlow, domain=DOMAIN): @callback def async_get_options_flow( config_entry: config_entries.ConfigEntry, - ) -> config_entries.OptionsFlow: + ) -> SpeedTestOptionsFlowHandler: """Get the options flow for this handler.""" return SpeedTestOptionsFlowHandler(config_entry) diff --git a/homeassistant/components/speedtestdotnet/const.py b/homeassistant/components/speedtestdotnet/const.py index 4843cf3f60d..c1979c133aa 100644 --- a/homeassistant/components/speedtestdotnet/const.py +++ b/homeassistant/components/speedtestdotnet/const.py @@ -1,52 +1,10 @@ """Constants used by Speedtest.net.""" from __future__ import annotations -from collections.abc import Callable -from dataclasses import dataclass from typing import Final -from homeassistant.components.sensor import ( - SensorDeviceClass, - SensorEntityDescription, - SensorStateClass, -) -from homeassistant.const import TIME_MILLISECONDS, Platform, UnitOfDataRate - DOMAIN: Final = "speedtestdotnet" - -@dataclass -class SpeedtestSensorEntityDescription(SensorEntityDescription): - """Class describing Speedtest sensor entities.""" - - value: Callable = round - - -SENSOR_TYPES: Final[tuple[SpeedtestSensorEntityDescription, ...]] = ( - SpeedtestSensorEntityDescription( - key="ping", - name="Ping", - native_unit_of_measurement=TIME_MILLISECONDS, - state_class=SensorStateClass.MEASUREMENT, - ), - SpeedtestSensorEntityDescription( - key="download", - name="Download", - device_class=SensorDeviceClass.DATA_RATE, - native_unit_of_measurement=UnitOfDataRate.MEGABITS_PER_SECOND, - state_class=SensorStateClass.MEASUREMENT, - value=lambda value: round(value / 10**6, 2), - ), - SpeedtestSensorEntityDescription( - key="upload", - name="Upload", - device_class=SensorDeviceClass.DATA_RATE, - native_unit_of_measurement=UnitOfDataRate.MEGABITS_PER_SECOND, - state_class=SensorStateClass.MEASUREMENT, - value=lambda value: round(value / 10**6, 2), - ), -) - CONF_SERVER_NAME: Final = "server_name" CONF_SERVER_ID: Final = "server_id" CONF_MANUAL: Final = "manual" @@ -65,5 +23,3 @@ DEFAULT_SERVER: Final = "*Auto Detect" ATTRIBUTION: Final = "Data retrieved from Speedtest.net by Ookla" ICON: Final = "mdi:speedometer" - -PLATFORMS: Final = [Platform.SENSOR] diff --git a/homeassistant/components/speedtestdotnet/coordinator.py b/homeassistant/components/speedtestdotnet/coordinator.py new file mode 100644 index 00000000000..4b088e00cf0 --- /dev/null +++ b/homeassistant/components/speedtestdotnet/coordinator.py @@ -0,0 +1,79 @@ +"""Coordinator for speedtestdotnet.""" + +import logging +from typing import Any + +import speedtest + +from homeassistant.config_entries import ConfigEntry +from homeassistant.core import HomeAssistant +from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed + +from .const import CONF_SERVER_ID, DEFAULT_SERVER, DOMAIN + +_LOGGER = logging.getLogger(__name__) + + +class SpeedTestDataCoordinator(DataUpdateCoordinator[dict[str, Any]]): + """Get the latest data from speedtest.net.""" + + config_entry: ConfigEntry + + def __init__( + self, hass: HomeAssistant, config_entry: ConfigEntry, api: speedtest.Speedtest + ) -> None: + """Initialize the data object.""" + self.hass = hass + self.config_entry = config_entry + self.api = api + self.servers: dict[str, dict] = {DEFAULT_SERVER: {}} + super().__init__( + self.hass, + _LOGGER, + name=DOMAIN, + ) + + def update_servers(self) -> None: + """Update list of test servers.""" + test_servers = self.api.get_servers() + test_servers_list = [] + for servers in test_servers.values(): + for server in servers: + test_servers_list.append(server) + for server in sorted( + test_servers_list, + key=lambda server: ( + server["country"], + server["name"], + server["sponsor"], + ), + ): + self.servers[ + f"{server['country']} - {server['sponsor']} - {server['name']}" + ] = server + + def update_data(self) -> dict[str, Any]: + """Get the latest data from speedtest.net.""" + self.update_servers() + self.api.closest.clear() + if self.config_entry.options.get(CONF_SERVER_ID): + server_id = self.config_entry.options.get(CONF_SERVER_ID) + self.api.get_servers(servers=[server_id]) + + best_server = self.api.get_best_server() + _LOGGER.debug( + "Executing speedtest.net speed test with server_id: %s", + best_server["id"], + ) + self.api.download() + self.api.upload() + return self.api.results.dict() + + async def _async_update_data(self) -> dict[str, Any]: + """Update Speedtest data.""" + try: + return await self.hass.async_add_executor_job(self.update_data) + except speedtest.NoMatchedServers as err: + raise UpdateFailed("Selected server is not found.") from err + except speedtest.SpeedtestException as err: + raise UpdateFailed(err) from err diff --git a/homeassistant/components/speedtestdotnet/sensor.py b/homeassistant/components/speedtestdotnet/sensor.py index e1ba0f560bb..66cf1bb5810 100644 --- a/homeassistant/components/speedtestdotnet/sensor.py +++ b/homeassistant/components/speedtestdotnet/sensor.py @@ -1,10 +1,17 @@ """Support for Speedtest.net internet speed testing sensor.""" from __future__ import annotations +from collections.abc import Callable +from dataclasses import dataclass from typing import Any, cast -from homeassistant.components.sensor import SensorEntity +from homeassistant.components.sensor import ( + SensorEntity, + SensorEntityDescription, + SensorStateClass, +) from homeassistant.config_entries import ConfigEntry +from homeassistant.const import TIME_MILLISECONDS, UnitOfDataRate from homeassistant.core import HomeAssistant from homeassistant.helpers.device_registry import DeviceEntryType from homeassistant.helpers.entity import DeviceInfo @@ -13,7 +20,6 @@ from homeassistant.helpers.restore_state import RestoreEntity from homeassistant.helpers.typing import StateType from homeassistant.helpers.update_coordinator import CoordinatorEntity -from . import SpeedTestDataCoordinator from .const import ( ATTR_BYTES_RECEIVED, ATTR_BYTES_SENT, @@ -24,8 +30,38 @@ from .const import ( DEFAULT_NAME, DOMAIN, ICON, - SENSOR_TYPES, - SpeedtestSensorEntityDescription, +) +from .coordinator import SpeedTestDataCoordinator + + +@dataclass +class SpeedtestSensorEntityDescription(SensorEntityDescription): + """Class describing Speedtest sensor entities.""" + + value: Callable = round + + +SENSOR_TYPES: tuple[SpeedtestSensorEntityDescription, ...] = ( + SpeedtestSensorEntityDescription( + key="ping", + name="Ping", + native_unit_of_measurement=TIME_MILLISECONDS, + state_class=SensorStateClass.MEASUREMENT, + ), + SpeedtestSensorEntityDescription( + key="download", + name="Download", + native_unit_of_measurement=UnitOfDataRate.MEGABITS_PER_SECOND, + state_class=SensorStateClass.MEASUREMENT, + value=lambda value: round(value / 10**6, 2), + ), + SpeedtestSensorEntityDescription( + key="upload", + name="Upload", + native_unit_of_measurement=UnitOfDataRate.MEGABITS_PER_SECOND, + state_class=SensorStateClass.MEASUREMENT, + value=lambda value: round(value / 10**6, 2), + ), ) diff --git a/tests/components/speedtestdotnet/test_sensor.py b/tests/components/speedtestdotnet/test_sensor.py index 06802a6cae7..9e05bb4c43b 100644 --- a/tests/components/speedtestdotnet/test_sensor.py +++ b/tests/components/speedtestdotnet/test_sensor.py @@ -1,13 +1,10 @@ """Tests for SpeedTest sensors.""" from unittest.mock import MagicMock -from homeassistant.components import speedtestdotnet from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN -from homeassistant.components.speedtestdotnet.const import ( - CONF_MANUAL, - DEFAULT_NAME, - SENSOR_TYPES, -) +from homeassistant.components.speedtestdotnet import DOMAIN +from homeassistant.components.speedtestdotnet.const import CONF_MANUAL, DEFAULT_NAME +from homeassistant.components.speedtestdotnet.sensor import SENSOR_TYPES from homeassistant.core import HomeAssistant, State from . import MOCK_RESULTS, MOCK_SERVERS, MOCK_STATES @@ -19,7 +16,7 @@ async def test_speedtestdotnet_sensors( hass: HomeAssistant, mock_api: MagicMock ) -> None: """Test sensors created for speedtestdotnet integration.""" - entry = MockConfigEntry(domain=speedtestdotnet.DOMAIN, data={}) + entry = MockConfigEntry(domain=DOMAIN, data={}) entry.add_to_hass(hass) mock_api.return_value.get_best_server.return_value = MOCK_SERVERS[1][0] @@ -45,9 +42,7 @@ async def test_restore_last_state(hass: HomeAssistant, mock_api: MagicMock) -> N for sensor, state in MOCK_STATES.items() ], ) - entry = MockConfigEntry( - domain=speedtestdotnet.DOMAIN, data={}, options={CONF_MANUAL: True} - ) + entry = MockConfigEntry(domain=DOMAIN, data={}, options={CONF_MANUAL: True}) entry.add_to_hass(hass) await hass.config_entries.async_setup(entry.entry_id)