Add DataUpdateCoordinator to pyLoad integration (#120237)

* Add DataUpdateCoordinator

* Update tests

* changes

* changes

* test coverage

* some changes

* Update homeassistant/components/pyload/sensor.py

* use dataclass

* fix ConfigEntry

* fix configtype

* fix some issues

* remove logger

* remove unnecessary else

* revert fixture changes

---------

Co-authored-by: Joost Lekkerkerker <joostlek@outlook.com>

homeassistant/components/pyload/__init__.py

@@ -20,9 +20,11 @@ from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryError, ConfigEntryNotReady
 from homeassistant.helpers.aiohttp_client import async_create_clientsession
 
+from .coordinator import PyLoadCoordinator
+
 PLATFORMS: list[Platform] = [Platform.SENSOR]
 
-type PyLoadConfigEntry = ConfigEntry[PyLoadAPI]
+type PyLoadConfigEntry = ConfigEntry[PyLoadCoordinator]
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: PyLoadConfigEntry) -> bool:
@@ -57,9 +59,11 @@ async def async_setup_entry(hass: HomeAssistant, entry: PyLoadConfigEntry) -> bo
         raise ConfigEntryError(
             f"Authentication failed for {entry.data[CONF_USERNAME]}, check your login credentials"
         ) from e
+    coordinator = PyLoadCoordinator(hass, pyloadapi)
 
-    entry.runtime_data = pyloadapi
+    await coordinator.async_config_entry_first_refresh()
 
+    entry.runtime_data = coordinator
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     return True

homeassistant/components/pyload/coordinator.py

@@ -0,0 +1,78 @@
+"""Update coordinator for pyLoad Integration."""
+
+from dataclasses import dataclass
+from datetime import timedelta
+import logging
+
+from pyloadapi import CannotConnect, InvalidAuth, ParserError, PyLoadAPI
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import ConfigEntryError
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+
+from .const import DOMAIN
+
+_LOGGER = logging.getLogger(__name__)
+
+SCAN_INTERVAL = timedelta(seconds=20)
+
+
+@dataclass(kw_only=True)
+class pyLoadData:
+    """Data from pyLoad."""
+
+    pause: bool
+    active: int
+    queue: int
+    total: int
+    speed: float
+    download: bool
+    reconnect: bool
+    captcha: bool
+    free_space: int
+
+
+class PyLoadCoordinator(DataUpdateCoordinator[pyLoadData]):
+    """pyLoad coordinator."""
+
+    config_entry: ConfigEntry
+
+    def __init__(self, hass: HomeAssistant, pyload: PyLoadAPI) -> None:
+        """Initialize pyLoad coordinator."""
+        super().__init__(
+            hass,
+            _LOGGER,
+            name=DOMAIN,
+            update_interval=SCAN_INTERVAL,
+        )
+        self.pyload = pyload
+        self.version: str | None = None
+
+    async def _async_update_data(self) -> pyLoadData:
+        """Fetch data from API endpoint."""
+        try:
+            if not self.version:
+                self.version = await self.pyload.version()
+            return pyLoadData(
+                **await self.pyload.get_status(),
+                free_space=await self.pyload.free_space(),
+            )
+
+        except InvalidAuth as e:
+            try:
+                await self.pyload.login()
+            except InvalidAuth as exc:
+                raise ConfigEntryError(
+                    f"Authentication failed for {self.pyload.username}, check your login credentials",
+                ) from exc
+
+            raise UpdateFailed(
+                "Unable to retrieve data due to cookie expiration but re-authentication was successful."
+            ) from e
+        except CannotConnect as e:
+            raise UpdateFailed(
+                "Unable to connect and retrieve data from pyLoad API"
+            ) from e
+        except ParserError as e:
+            raise UpdateFailed("Unable to parse data from pyLoad API") from e

homeassistant/components/pyload/sensor.py

@@ -2,18 +2,8 @@
 
 from __future__ import annotations
 
-from datetime import timedelta
 from enum import StrEnum
-import logging
-from time import monotonic
 
-from pyloadapi import (
-    CannotConnect,
-    InvalidAuth,
-    ParserError,
-    PyLoadAPI,
-    StatusServerResponse,
-)
 import voluptuous as vol
 
 from homeassistant.components.sensor import (
@@ -40,13 +30,11 @@ from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.issue_registry import IssueSeverity, async_create_issue
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType, StateType
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from . import PyLoadConfigEntry
 from .const import DEFAULT_HOST, DEFAULT_NAME, DEFAULT_PORT, DOMAIN, ISSUE_PLACEHOLDER
-
+from .coordinator import PyLoadCoordinator
-_LOGGER = logging.getLogger(__name__)
-
-SCAN_INTERVAL = timedelta(seconds=15)
 
 
 class PyLoadSensorEntity(StrEnum):
@@ -92,7 +80,6 @@ async def async_setup_platform(
     result = await hass.config_entries.flow.async_init(
         DOMAIN, context={"source": SOURCE_IMPORT}, data=config
     )
-    _LOGGER.debug(result)
     if (
         result.get("type") == FlowResultType.CREATE_ENTRY
         or result.get("reason") == "already_configured"
@@ -132,91 +119,45 @@ async def async_setup_entry(
 ) -> None:
     """Set up the pyLoad sensors."""
 
-    pyloadapi = entry.runtime_data
+    coordinator = entry.runtime_data
 
     async_add_entities(
         (
             PyLoadSensor(
-                api=pyloadapi,
+                coordinator=coordinator,
                 entity_description=description,
-                client_name=entry.title,
-                entry_id=entry.entry_id,
             )
             for description in SENSOR_DESCRIPTIONS
         ),
-        True,
     )
 
 
-class PyLoadSensor(SensorEntity):
+class PyLoadSensor(CoordinatorEntity[PyLoadCoordinator], SensorEntity):
     """Representation of a pyLoad sensor."""
 
     _attr_has_entity_name = True
 
     def __init__(
         self,
-        api: PyLoadAPI,
+        coordinator: PyLoadCoordinator,
         entity_description: SensorEntityDescription,
-        client_name: str,
-        entry_id: str,
     ) -> None:
         """Initialize a new pyLoad sensor."""
-        self.type = entity_description.key
+        super().__init__(coordinator)
-        self.api = api
+        self._attr_unique_id = (
-        self._attr_unique_id = f"{entry_id}_{entity_description.key}"
+            f"{coordinator.config_entry.entry_id}_{entity_description.key}"
+        )
         self.entity_description = entity_description
-        self._attr_available = False
-        self.data: StatusServerResponse
         self.device_info = DeviceInfo(
             entry_type=DeviceEntryType.SERVICE,
             manufacturer="PyLoad Team",
             model="pyLoad",
-            configuration_url=api.api_url,
+            configuration_url=coordinator.pyload.api_url,
-            identifiers={(DOMAIN, entry_id)},
+            identifiers={(DOMAIN, coordinator.config_entry.entry_id)},
+            sw_version=coordinator.version,
         )
 
-    async def async_update(self) -> None:
-        """Update state of sensor."""
-        start = monotonic()
-        try:
-            status = await self.api.get_status()
-        except InvalidAuth:
-            _LOGGER.info("Authentication failed, trying to reauthenticate")
-            try:
-                await self.api.login()
-            except InvalidAuth:
-                _LOGGER.error(
-                    "Authentication failed for %s, check your login credentials",
-                    self.api.username,
-                )
-                return
-            else:
-                _LOGGER.info(
-                    "Unable to retrieve data due to cookie expiration "
-                    "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
-            _LOGGER.debug(
-                "Finished fetching pyload data in %.3f seconds",
-                monotonic() - start,
-            )
-
-        self._attr_available = True
-
     @property
     def native_value(self) -> StateType:
         """Return the state of the sensor."""
-        return self.data.get(self.entity_description.key)
+        return getattr(self.coordinator.data, self.entity_description.key)

tests/components/pyload/conftest.py

@@ -66,12 +66,10 @@ def mock_pyloadapi() -> Generator[AsyncMock, None, None]:
             "homeassistant.components.pyload.PyLoadAPI", autospec=True
         ) as mock_client,
         patch("homeassistant.components.pyload.config_flow.PyLoadAPI", new=mock_client),
-        patch("homeassistant.components.pyload.sensor.PyLoadAPI", new=mock_client),
     ):
         client = mock_client.return_value
         client.username = "username"
         client.api_url = "https://pyload.local:8000/"
-
         client.login.return_value = LoginResponse(
             {
                 "_permanent": True,
@@ -97,7 +95,7 @@ def mock_pyloadapi() -> Generator[AsyncMock, None, None]:
                 "captcha": False,
             }
         )
-
+        client.version.return_value = "0.5.0"
         client.free_space.return_value = 99999999999
         yield client
 

tests/components/pyload/snapshots/test_sensor.ambr

@@ -161,183 +161,6 @@
     'state': 'unavailable',
   })
 # ---
-# name: test_sensors_unavailable[CannotConnect][sensor.pyload_speed-entry]
-  EntityRegistryEntrySnapshot({
-    'aliases': set({
-    }),
-    'area_id': None,
-    'capabilities': None,
-    'config_entry_id': <ANY>,
-    'device_class': None,
-    'device_id': <ANY>,
-    'disabled_by': None,
-    'domain': 'sensor',
-    'entity_category': None,
-    'entity_id': 'sensor.pyload_speed',
-    'has_entity_name': True,
-    'hidden_by': None,
-    'icon': None,
-    'id': <ANY>,
-    'labels': set({
-    }),
-    'name': None,
-    'options': dict({
-      'sensor': dict({
-        'suggested_display_precision': 1,
-      }),
-      'sensor.private': dict({
-        'suggested_unit_of_measurement': <UnitOfDataRate.MEGABYTES_PER_SECOND: 'MB/s'>,
-      }),
-    }),
-    'original_device_class': <SensorDeviceClass.DATA_RATE: 'data_rate'>,
-    'original_icon': None,
-    'original_name': 'Speed',
-    'platform': 'pyload',
-    'previous_unique_id': None,
-    'supported_features': 0,
-    'translation_key': None,
-    'unique_id': 'XXXXXXXXXXXXXX_speed',
-    'unit_of_measurement': <UnitOfDataRate.MEGABYTES_PER_SECOND: 'MB/s'>,
-  })
-# ---
-# name: test_sensors_unavailable[CannotConnect][sensor.pyload_speed-state]
-  StateSnapshot({
-    'attributes': ReadOnlyDict({
-      'device_class': 'data_rate',
-      'friendly_name': 'pyLoad Speed',
-      'unit_of_measurement': <UnitOfDataRate.MEGABYTES_PER_SECOND: 'MB/s'>,
-    }),
-    'context': <ANY>,
-    'entity_id': 'sensor.pyload_speed',
-    'last_changed': <ANY>,
-    'last_reported': <ANY>,
-    'last_updated': <ANY>,
-    'state': 'unavailable',
-  })
-# ---
-# name: test_sensors_unavailable[InvalidAuth][sensor.pyload_speed-entry]
-  EntityRegistryEntrySnapshot({
-    'aliases': set({
-    }),
-    'area_id': None,
-    'capabilities': None,
-    'config_entry_id': <ANY>,
-    'device_class': None,
-    'device_id': <ANY>,
-    'disabled_by': None,
-    'domain': 'sensor',
-    'entity_category': None,
-    'entity_id': 'sensor.pyload_speed',
-    'has_entity_name': True,
-    'hidden_by': None,
-    'icon': None,
-    'id': <ANY>,
-    'labels': set({
-    }),
-    'name': None,
-    'options': dict({
-      'sensor': dict({
-        'suggested_display_precision': 1,
-      }),
-      'sensor.private': dict({
-        'suggested_unit_of_measurement': <UnitOfDataRate.MEGABYTES_PER_SECOND: 'MB/s'>,
-      }),
-    }),
-    'original_device_class': <SensorDeviceClass.DATA_RATE: 'data_rate'>,
-    'original_icon': None,
-    'original_name': 'Speed',
-    'platform': 'pyload',
-    'previous_unique_id': None,
-    'supported_features': 0,
-    'translation_key': None,
-    'unique_id': 'XXXXXXXXXXXXXX_speed',
-    'unit_of_measurement': <UnitOfDataRate.MEGABYTES_PER_SECOND: 'MB/s'>,
-  })
-# ---
-# name: test_sensors_unavailable[InvalidAuth][sensor.pyload_speed-state]
-  StateSnapshot({
-    'attributes': ReadOnlyDict({
-      'device_class': 'data_rate',
-      'friendly_name': 'pyLoad Speed',
-      'unit_of_measurement': <UnitOfDataRate.MEGABYTES_PER_SECOND: 'MB/s'>,
-    }),
-    'context': <ANY>,
-    'entity_id': 'sensor.pyload_speed',
-    'last_changed': <ANY>,
-    'last_reported': <ANY>,
-    'last_updated': <ANY>,
-    'state': 'unavailable',
-  })
-# ---
-# name: test_sensors_unavailable[ParserError][sensor.pyload_speed-entry]
-  EntityRegistryEntrySnapshot({
-    'aliases': set({
-    }),
-    'area_id': None,
-    'capabilities': None,
-    'config_entry_id': <ANY>,
-    'device_class': None,
-    'device_id': <ANY>,
-    'disabled_by': None,
-    'domain': 'sensor',
-    'entity_category': None,
-    'entity_id': 'sensor.pyload_speed',
-    'has_entity_name': True,
-    'hidden_by': None,
-    'icon': None,
-    'id': <ANY>,
-    'labels': set({
-    }),
-    'name': None,
-    'options': dict({
-      'sensor': dict({
-        'suggested_display_precision': 1,
-      }),
-      'sensor.private': dict({
-        'suggested_unit_of_measurement': <UnitOfDataRate.MEGABYTES_PER_SECOND: 'MB/s'>,
-      }),
-    }),
-    'original_device_class': <SensorDeviceClass.DATA_RATE: 'data_rate'>,
-    'original_icon': None,
-    'original_name': 'Speed',
-    'platform': 'pyload',
-    'previous_unique_id': None,
-    'supported_features': 0,
-    'translation_key': None,
-    'unique_id': 'XXXXXXXXXXXXXX_speed',
-    'unit_of_measurement': <UnitOfDataRate.MEGABYTES_PER_SECOND: 'MB/s'>,
-  })
-# ---
-# name: test_sensors_unavailable[ParserError][sensor.pyload_speed-state]
-  StateSnapshot({
-    'attributes': ReadOnlyDict({
-      'device_class': 'data_rate',
-      'friendly_name': 'pyLoad Speed',
-      'unit_of_measurement': <UnitOfDataRate.MEGABYTES_PER_SECOND: 'MB/s'>,
-    }),
-    'context': <ANY>,
-    'entity_id': 'sensor.pyload_speed',
-    'last_changed': <ANY>,
-    'last_reported': <ANY>,
-    'last_updated': <ANY>,
-    'state': 'unavailable',
-  })
-# ---
-# name: test_setup
-  StateSnapshot({
-    'attributes': ReadOnlyDict({
-      'device_class': 'data_rate',
-      'friendly_name': 'pyload Speed',
-      'unit_of_measurement': <UnitOfDataRate.MEGABYTES_PER_SECOND: 'MB/s'>,
-    }),
-    'context': <ANY>,
-    'entity_id': 'sensor.pyload_speed',
-    'last_changed': <ANY>,
-    'last_reported': <ANY>,
-    'last_updated': <ANY>,
-    'state': '5.405963',
-  })
-# ---
 # name: test_setup[sensor.pyload_speed-entry]
   EntityRegistryEntrySnapshot({
     'aliases': set({

tests/components/pyload/test_sensor.py

@@ -8,7 +8,7 @@ import pytest
 from syrupy.assertion import SnapshotAssertion
 
 from homeassistant.components.pyload.const import DOMAIN
-from homeassistant.components.pyload.sensor import SCAN_INTERVAL
+from homeassistant.components.pyload.coordinator import SCAN_INTERVAL
 from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import DOMAIN as HOMEASSISTANT_DOMAIN, HomeAssistant