Refactor sensor platform of Pyload integration (#119716)

homeassistant/components/pyload/sensor.py

@@ -3,12 +3,14 @@
 from __future__ import annotations
 
 from datetime import timedelta
+from enum import StrEnum
 import logging
+from time import monotonic
+from typing import Any
 
 from aiohttp import CookieJar
 from pyloadapi.api import PyLoadAPI
 from pyloadapi.exceptions import CannotConnect, InvalidAuth, ParserError
-from pyloadapi.types import StatusServerResponse
 import voluptuous as vol
 
 from homeassistant.components.sensor import (
@@ -32,29 +34,37 @@ from homeassistant.exceptions import PlatformNotReady
 from homeassistant.helpers.aiohttp_client import async_create_clientsession
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
+from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType, StateType
 
 from .const import DEFAULT_HOST, DEFAULT_NAME, DEFAULT_PORT
 
 _LOGGER = logging.getLogger(__name__)
 
-
 SCAN_INTERVAL = timedelta(seconds=15)
 
-SENSOR_TYPES = {
-    "speed": SensorEntityDescription(
-        key="speed",
+
+class PyLoadSensorEntity(StrEnum):
+    """pyLoad Sensor Entities."""
+
+    SPEED = "speed"
+
+
+SENSOR_DESCRIPTIONS: tuple[SensorEntityDescription, ...] = (
+    SensorEntityDescription(
+        key=PyLoadSensorEntity.SPEED,
         name="Speed",
-        native_unit_of_measurement=UnitOfDataRate.MEGABYTES_PER_SECOND,
         device_class=SensorDeviceClass.DATA_RATE,
+        native_unit_of_measurement=UnitOfDataRate.BYTES_PER_SECOND,
+        suggested_unit_of_measurement=UnitOfDataRate.MEGABYTES_PER_SECOND,
+        suggested_display_precision=1,
+    ),
 )
-}
 
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
     {
         vol.Optional(CONF_HOST, default=DEFAULT_HOST): cv.string,
         vol.Optional(CONF_MONITORED_VARIABLES, default=["speed"]): vol.All(
-            cv.ensure_list, [vol.In(SENSOR_TYPES)]
+            cv.ensure_list, [vol.In(PyLoadSensorEntity)]
         ),
         vol.Optional(CONF_NAME, default=DEFAULT_NAME): cv.string,
         vol.Optional(CONF_PASSWORD): cv.string,
@@ -78,7 +88,6 @@ async def async_setup_platform(
     name = config[CONF_NAME]
     username = config.get(CONF_USERNAME)
     password = config.get(CONF_PASSWORD)
-    monitored_types = config[CONF_MONITORED_VARIABLES]
     url = f"{protocol}://{host}:{port}/"
 
     session = async_create_clientsession(
@@ -100,33 +109,36 @@ async def async_setup_platform(
             f"Authentication failed for {config[CONF_USERNAME]}, check your login credentials"
         ) from e
 
-    devices = []
-    for ng_type in monitored_types:
-        new_sensor = PyLoadSensor(
-            api=pyloadapi, sensor_type=SENSOR_TYPES[ng_type], client_name=name
+    async_add_entities(
+        (
+            PyLoadSensor(
+                api=pyloadapi, entity_description=description, client_name=name
+            )
+            for description in SENSOR_DESCRIPTIONS
+        ),
+        True,
     )
-        devices.append(new_sensor)
-
-    async_add_entities(devices, True)
 
 
 class PyLoadSensor(SensorEntity):
     """Representation of a pyLoad sensor."""
 
     def __init__(
-        self, api: PyLoadAPI, sensor_type: SensorEntityDescription, client_name
+        self, api: PyLoadAPI, entity_description: SensorEntityDescription, client_name
     ) -> None:
         """Initialize a new pyLoad sensor."""
-        self._attr_name = f"{client_name} {sensor_type.name}"
-        self.type = sensor_type.key
+        self._attr_name = f"{client_name} {entity_description.name}"
+        self.type = entity_description.key
         self.api = api
-        self.entity_description = sensor_type
-        self.data: StatusServerResponse
+        self.entity_description = entity_description
+        self._attr_available = False
+        self.data: dict[str, Any] = {}
 
     async def async_update(self) -> None:
         """Update state of sensor."""
+        start = monotonic()
         try:
-            self.data = await self.api.get_status()
+            status = await self.api.get_status()
         except InvalidAuth:
             _LOGGER.info("Authentication failed, trying to reauthenticate")
             try:
@@ -143,15 +155,27 @@ class PyLoadSensor(SensorEntity):
                     "but re-authentication was successful"
                 )
                 return
+            finally:
+                self._attr_available = False
+
         except CannotConnect:
             _LOGGER.debug("Unable to connect and retrieve data from pyLoad API")
+            self._attr_available = False
             return
         except ParserError:
             _LOGGER.error("Unable to parse data from pyLoad API")
+            self._attr_available = False
             return
+        else:
+            self.data = status.to_dict()
+            _LOGGER.debug(
+                "Finished fetching pyload data in %.3f seconds",
+                monotonic() - start,
+            )
 
-        value = getattr(self.data, self.type)
+        self._attr_available = True
 
-        if "speed" in self.type and value > 0:
-            # Convert download rate from Bytes/s to MBytes/s
-            self._attr_native_value = round(value / 2**20, 2)
+    @property
+    def native_value(self) -> StateType:
+        """Return the state of the sensor."""
+        return self.data.get(self.entity_description.key)

tests/components/pyload/conftest.py

@@ -71,4 +71,5 @@ def mock_pyloadapi() -> Generator[AsyncMock, None, None]:
                 "captcha": False,
             }
         )
+        client.free_space.return_value = 99999999999
         yield client

tests/components/pyload/snapshots/test_sensor.ambr

@@ -11,6 +11,6 @@
     'last_changed': <ANY>,
     'last_reported': <ANY>,
     'last_updated': <ANY>,
-    'state': '5.16',
+    'state': '5.405963',
   })
 # ---

tests/components/pyload/test_sensor.py

@@ -9,12 +9,15 @@ from syrupy.assertion import SnapshotAssertion
 
 from homeassistant.components.pyload.sensor import SCAN_INTERVAL
 from homeassistant.components.sensor import DOMAIN
+from homeassistant.const import STATE_UNAVAILABLE
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.typing import ConfigType
 from homeassistant.setup import async_setup_component
 
 from tests.common import async_fire_time_changed
 
+SENSORS = ["sensor.pyload_speed"]
+
 
 @pytest.mark.usefixtures("mock_pyloadapi")
 async def test_setup(
@@ -27,7 +30,8 @@ async def test_setup(
     assert await async_setup_component(hass, DOMAIN, pyload_config)
     await hass.async_block_till_done()
 
-    result = hass.states.get("sensor.pyload_speed")
+    for sensor in SENSORS:
+        result = hass.states.get(sensor)
         assert result == snapshot
 
 
@@ -76,6 +80,8 @@ async def test_sensor_update_exceptions(
     exception: Exception,
     expected_exception: str,
     caplog: pytest.LogCaptureFixture,
+    snapshot: SnapshotAssertion,
+    freezer: FrozenDateTimeFactory,
 ) -> None:
     """Test exceptions during update of pyLoad sensor."""
 
@@ -87,6 +93,9 @@ async def test_sensor_update_exceptions(
     assert len(hass.states.async_all(DOMAIN)) == 1
     assert expected_exception in caplog.text
 
+    for sensor in SENSORS:
+        assert hass.states.get(sensor).state == STATE_UNAVAILABLE
+
 
 async def test_sensor_invalid_auth(
     hass: HomeAssistant,