Add device_info to Speedtestdotnet and some code cleanup (#56612)

* Apply code cleanup suggestions from previous PRs

* Update homeassistant/components/speedtestdotnet/const.py

Co-authored-by: Franck Nijhof <git@frenck.dev>

* fix native_value, and ping value in test

* use self._state instead of _attr_native_value

* update identifiers and add more tests

Co-authored-by: Franck Nijhof <git@frenck.dev>

homeassistant/components/speedtestdotnet/__init__.py

@@ -83,6 +83,11 @@ class SpeedTestDataCoordinator(DataUpdateCoordinator):
             update_method=self.async_update,
         )
 
+    def initialize(self) -> None:
+        """Initialize speedtest api."""
+        self.api = speedtest.Speedtest()
+        self.update_servers()
+
     def update_servers(self):
         """Update list of test servers."""
         test_servers = self.api.get_servers()
@@ -131,8 +136,7 @@ class SpeedTestDataCoordinator(DataUpdateCoordinator):
     async def async_setup(self) -> None:
         """Set up SpeedTest."""
         try:
-            self.api = await self.hass.async_add_executor_job(speedtest.Speedtest)
-            await self.hass.async_add_executor_job(self.update_servers)
+            await self.hass.async_add_executor_job(self.initialize)
         except speedtest.SpeedtestException as err:
             raise ConfigEntryNotReady from err
 

homeassistant/components/speedtestdotnet/const.py

@@ -1,7 +1,8 @@
-"""Consts used by Speedtest.net."""
+"""Constants used by Speedtest.net."""
 from __future__ import annotations
 
-from typing import Final
+from dataclasses import dataclass
+from typing import Callable, Final
 
 from homeassistant.components.sensor import (
     STATE_CLASS_MEASUREMENT,
@@ -13,24 +14,34 @@ DOMAIN: Final = "speedtestdotnet"
 
 SPEED_TEST_SERVICE: Final = "speedtest"
 
-SENSOR_TYPES: Final[tuple[SensorEntityDescription, ...]] = (
-    SensorEntityDescription(
+
+@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=STATE_CLASS_MEASUREMENT,
     ),
-    SensorEntityDescription(
+    SpeedtestSensorEntityDescription(
         key="download",
         name="Download",
         native_unit_of_measurement=DATA_RATE_MEGABITS_PER_SECOND,
         state_class=STATE_CLASS_MEASUREMENT,
+        value=lambda value: round(value / 10 ** 6, 2),
     ),
-    SensorEntityDescription(
+    SpeedtestSensorEntityDescription(
         key="upload",
         name="Upload",
         native_unit_of_measurement=DATA_RATE_MEGABITS_PER_SECOND,
         state_class=STATE_CLASS_MEASUREMENT,
+        value=lambda value: round(value / 10 ** 6, 2),
     ),
 )
 

homeassistant/components/speedtestdotnet/sensor.py

@@ -1,15 +1,16 @@
 """Support for Speedtest.net internet speed testing sensor."""
 from __future__ import annotations
 
-from typing import Any
+from typing import Any, cast
 
-from homeassistant.components.sensor import SensorEntity, SensorEntityDescription
+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 HomeAssistant, callback
+from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.restore_state import RestoreEntity
+from homeassistant.helpers.typing import StateType
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from .const import (
@@ -23,6 +24,7 @@ from .const import (
     DOMAIN,
     ICON,
     SENSOR_TYPES,
+    SpeedtestSensorEntityDescription,
 )
 
 
@@ -43,19 +45,34 @@ class SpeedtestSensor(CoordinatorEntity, RestoreEntity, SensorEntity):
     """Implementation of a speedtest.net sensor."""
 
     coordinator: SpeedTestDataCoordinator
+    entity_description: SpeedtestSensorEntityDescription
     _attr_icon = ICON
 
     def __init__(
         self,
         coordinator: SpeedTestDataCoordinator,
-        description: SensorEntityDescription,
+        description: SpeedtestSensorEntityDescription,
     ) -> None:
         """Initialize the sensor."""
         super().__init__(coordinator)
         self.entity_description = description
         self._attr_name = f"{DEFAULT_NAME} {description.name}"
         self._attr_unique_id = description.key
+        self._state: StateType = None
         self._attrs = {ATTR_ATTRIBUTION: ATTRIBUTION}
+        self._attr_device_info = {
+            "identifiers": {(DOMAIN, self.coordinator.config_entry.entry_id)},
+            "name": DEFAULT_NAME,
+            "entry_type": "service",
+        }
+
+    @property
+    def native_value(self) -> StateType:
+        """Return native value for entity."""
+        if self.coordinator.data:
+            state = self.coordinator.data[self.entity_description.key]
+            self._state = cast(StateType, self.entity_description.value(state))
+        return self._state
 
     @property
     def extra_state_attributes(self) -> dict[str, Any]:
@@ -83,27 +100,4 @@ class SpeedtestSensor(CoordinatorEntity, RestoreEntity, SensorEntity):
         await super().async_added_to_hass()
         state = await self.async_get_last_state()
         if state:
-            self._attr_native_value = state.state
-
-        @callback
-        def update() -> None:
-            """Update state."""
-            self._update_state()
-            self.async_write_ha_state()
-
-        self.async_on_remove(self.coordinator.async_add_listener(update))
-        self._update_state()
-
-    def _update_state(self):
-        """Update sensors state."""
-        if self.coordinator.data:
-            if self.entity_description.key == "ping":
-                self._attr_native_value = self.coordinator.data["ping"]
-            elif self.entity_description.key == "download":
-                self._attr_native_value = round(
-                    self.coordinator.data["download"] / 10 ** 6, 2
-                )
-            elif self.entity_description.key == "upload":
-                self._attr_native_value = round(
-                    self.coordinator.data["upload"] / 10 ** 6, 2
-                )
+            self._state = state.state

tests/components/speedtestdotnet/__init__.py

@@ -52,4 +52,4 @@ MOCK_RESULTS = {
     "share": None,
 }
 
-MOCK_STATES = {"ping": "18.465", "download": "1.02", "upload": "1.02"}
+MOCK_STATES = {"ping": "18", "download": "1.02", "upload": "1.02"}

tests/components/speedtestdotnet/test_config_flow.py

@@ -65,6 +65,28 @@ async def test_options(hass: HomeAssistant, mock_api: MagicMock) -> None:
 
     assert hass.data[DOMAIN].update_interval is None
 
+    # test setting server name to "*Auto Detect"
+    result = await hass.config_entries.options.async_init(entry.entry_id)
+    assert result["type"] == data_entry_flow.RESULT_TYPE_FORM
+    assert result["step_id"] == "init"
+
+    result = await hass.config_entries.options.async_configure(
+        result["flow_id"],
+        user_input={
+            CONF_SERVER_NAME: "*Auto Detect",
+            CONF_SCAN_INTERVAL: 30,
+            CONF_MANUAL: True,
+        },
+    )
+
+    assert result["type"] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY
+    assert result["data"] == {
+        CONF_SERVER_NAME: "*Auto Detect",
+        CONF_SERVER_ID: None,
+        CONF_SCAN_INTERVAL: 30,
+        CONF_MANUAL: True,
+    }
+
     # test setting the option to update periodically
     result2 = await hass.config_entries.options.async_init(entry.entry_id)
     assert result2["type"] == data_entry_flow.RESULT_TYPE_FORM

tests/components/speedtestdotnet/test_init.py

@@ -1,4 +1,5 @@
 """Tests for SpeedTest integration."""
+from datetime import timedelta
 from unittest.mock import MagicMock
 
 import speedtest
@@ -13,8 +14,9 @@ from homeassistant.components.speedtestdotnet.const import (
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.const import CONF_SCAN_INTERVAL, STATE_UNAVAILABLE
 from homeassistant.core import HomeAssistant
+import homeassistant.util.dt as dt_util
 
-from tests.common import MockConfigEntry
+from tests.common import MockConfigEntry, async_fire_time_changed
 
 
 async def test_successful_config_entry(hass: HomeAssistant) -> None:
@@ -74,6 +76,10 @@ async def test_server_not_found(hass: HomeAssistant, mock_api: MagicMock) -> Non
 
     entry = MockConfigEntry(
         domain=DOMAIN,
+        options={
+            CONF_MANUAL: False,
+            CONF_SCAN_INTERVAL: 60,
+        },
     )
     entry.add_to_hass(hass)
 
@@ -82,7 +88,10 @@ async def test_server_not_found(hass: HomeAssistant, mock_api: MagicMock) -> Non
     assert hass.data[DOMAIN]
 
     mock_api.return_value.get_servers.side_effect = speedtest.NoMatchedServers
-    await hass.data[DOMAIN].async_refresh()
+    async_fire_time_changed(
+        hass,
+        dt_util.utcnow() + timedelta(minutes=entry.options[CONF_SCAN_INTERVAL] + 1),
+    )
     await hass.async_block_till_done()
     state = hass.states.get("sensor.speedtest_ping")
     assert state is not None

tests/components/speedtestdotnet/test_sensor.py

@@ -3,12 +3,16 @@ 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 DEFAULT_NAME, SENSOR_TYPES
-from homeassistant.core import HomeAssistant
+from homeassistant.components.speedtestdotnet.const import (
+    CONF_MANUAL,
+    DEFAULT_NAME,
+    SENSOR_TYPES,
+)
+from homeassistant.core import HomeAssistant, State
 
 from . import MOCK_RESULTS, MOCK_SERVERS, MOCK_STATES
 
-from tests.common import MockConfigEntry
+from tests.common import MockConfigEntry, mock_restore_cache
 
 
 async def test_speedtestdotnet_sensors(
@@ -30,3 +34,28 @@ async def test_speedtestdotnet_sensors(
         sensor = hass.states.get(f"sensor.{DEFAULT_NAME}_{description.name}")
         assert sensor
         assert sensor.state == MOCK_STATES[description.key]
+
+
+async def test_restore_last_state(hass: HomeAssistant, mock_api: MagicMock) -> None:
+    """Test restoring last state for sensors."""
+    mock_restore_cache(
+        hass,
+        [
+            State(f"sensor.speedtest_{sensor}", state)
+            for sensor, state in MOCK_STATES.items()
+        ],
+    )
+    entry = MockConfigEntry(
+        domain=speedtestdotnet.DOMAIN, data={}, options={CONF_MANUAL: True}
+    )
+    entry.add_to_hass(hass)
+
+    await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 3
+
+    for description in SENSOR_TYPES:
+        sensor = hass.states.get(f"sensor.speedtest_{description.name}")
+        assert sensor
+        assert sensor.state == MOCK_STATES[description.key]