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.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.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
-
-_LOGGER = logging.getLogger(__name__)
-
-SCAN_INTERVAL = timedelta(seconds=15)
+from .coordinator import PyLoadCoordinator
 
 
 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
-        self.api = api
-        self._attr_unique_id = f"{entry_id}_{entity_description.key}"
+        super().__init__(coordinator)
+        self._attr_unique_id = (
+            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,
-            identifiers={(DOMAIN, entry_id)},
+            configuration_url=coordinator.pyload.api_url,
+            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