Subclass a DataUpdateCoordinator for Ridwell (home-assistant#85644)

.coveragerc

@@ -1073,6 +1073,7 @@ omit =
     homeassistant/components/rest/switch.py
     homeassistant/components/rfxtrx/diagnostics.py
     homeassistant/components/ridwell/__init__.py
+    homeassistant/components/ridwell/coordinator.py
     homeassistant/components/ridwell/entity.py
     homeassistant/components/ridwell/sensor.py
     homeassistant/components/ridwell/switch.py

homeassistant/components/ridwell/__init__.py

@@ -1,84 +1,24 @@
 """The Ridwell integration."""
 from __future__ import annotations
 
-import asyncio
-from dataclasses import dataclass
-from datetime import timedelta
 from typing import Any
 
-from aioridwell import async_get_client
-from aioridwell.errors import InvalidCredentialsError, RidwellError
-from aioridwell.model import RidwellAccount, RidwellPickupEvent
-
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_PASSWORD, CONF_USERNAME, Platform
+from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
-from homeassistant.helpers import aiohttp_client, entity_registry as er
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+from homeassistant.helpers import entity_registry as er
 
 from .const import DOMAIN, LOGGER, SENSOR_TYPE_NEXT_PICKUP
-
-DEFAULT_UPDATE_INTERVAL = timedelta(hours=1)
+from .coordinator import RidwellDataUpdateCoordinator
 
 PLATFORMS: list[Platform] = [Platform.SENSOR, Platform.SWITCH]
 
 
-@dataclass
-class RidwellData:
-    """Define an object to be stored in `hass.data`."""
-
-    accounts: dict[str, RidwellAccount]
-    coordinator: DataUpdateCoordinator[dict[str, RidwellPickupEvent]]
-
-
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Ridwell from a config entry."""
-    session = aiohttp_client.async_get_clientsession(hass)
-
-    try:
-        client = await async_get_client(
-            entry.data[CONF_USERNAME], entry.data[CONF_PASSWORD], session=session
-        )
-    except InvalidCredentialsError as err:
-        raise ConfigEntryAuthFailed("Invalid username/password") from err
-    except RidwellError as err:
-        raise ConfigEntryNotReady(err) from err
-
-    accounts = await client.async_get_accounts()
-
-    async def async_update_data() -> dict[str, RidwellPickupEvent]:
-        """Get the latest pickup events."""
-        data = {}
-
-        async def async_get_pickups(account: RidwellAccount) -> None:
-            """Get the latest pickups for an account."""
-            data[account.account_id] = await account.async_get_next_pickup_event()
-
-        tasks = [async_get_pickups(account) for account in accounts.values()]
-        results = await asyncio.gather(*tasks, return_exceptions=True)
-        for result in results:
-            if isinstance(result, InvalidCredentialsError):
-                raise ConfigEntryAuthFailed("Invalid username/password") from result
-            if isinstance(result, RidwellError):
-                raise UpdateFailed(result) from result
-
-        return data
-
-    coordinator: DataUpdateCoordinator[
-        dict[str, RidwellPickupEvent]
-    ] = DataUpdateCoordinator(
-        hass,
-        LOGGER,
-        name=entry.title,
-        update_interval=DEFAULT_UPDATE_INTERVAL,
-        update_method=async_update_data,
-    )
-
-    await coordinator.async_config_entry_first_refresh()
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = RidwellData(
-        accounts=accounts, coordinator=coordinator
-    )
+    coordinator = RidwellDataUpdateCoordinator(hass, name=entry.title)
+    await coordinator.async_initialize()
+    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 

homeassistant/components/ridwell/coordinator.py

@@ -0,0 +1,78 @@
+"""Define a Ridwell coordinator."""
+from __future__ import annotations
+
+import asyncio
+from datetime import timedelta
+from typing import cast
+
+from aioridwell.client import async_get_client
+from aioridwell.errors import InvalidCredentialsError, RidwellError
+from aioridwell.model import RidwellAccount, RidwellPickupEvent
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
+from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
+from homeassistant.helpers import aiohttp_client
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+
+from .const import LOGGER
+
+UPDATE_INTERVAL = timedelta(hours=1)
+
+
+class RidwellDataUpdateCoordinator(
+    DataUpdateCoordinator[dict[str, RidwellPickupEvent]]
+):
+    """Class to manage fetching data from single endpoint."""
+
+    config_entry: ConfigEntry
+
+    def __init__(self, hass: HomeAssistant, *, name: str) -> None:
+        """Initialize global data updater."""
+        # These will be filled in by async_initialize; we give them these defaults to
+        # avoid arduous typing checks down the line:
+        self.accounts: dict[str, RidwellAccount] = {}
+        self.dashboard_url = ""
+        self.user_id = ""
+
+        super().__init__(hass, LOGGER, name=name, update_interval=UPDATE_INTERVAL)
+
+    async def _async_update_data(self) -> dict[str, RidwellPickupEvent]:
+        """Fetch the latest data from the source."""
+        data = {}
+
+        async def async_get_pickups(account: RidwellAccount) -> None:
+            """Get the latest pickups for an account."""
+            data[account.account_id] = await account.async_get_next_pickup_event()
+
+        tasks = [async_get_pickups(account) for account in self.accounts.values()]
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+        for result in results:
+            if isinstance(result, InvalidCredentialsError):
+                raise ConfigEntryAuthFailed("Invalid username/password") from result
+            if isinstance(result, RidwellError):
+                raise UpdateFailed(result) from result
+
+        return data
+
+    async def async_initialize(self) -> None:
+        """Initialize the coordinator."""
+        session = aiohttp_client.async_get_clientsession(self.hass)
+
+        try:
+            client = await async_get_client(
+                self.config_entry.data[CONF_USERNAME],
+                self.config_entry.data[CONF_PASSWORD],
+                session=session,
+            )
+        except InvalidCredentialsError as err:
+            raise ConfigEntryAuthFailed("Invalid username/password") from err
+        except RidwellError as err:
+            raise ConfigEntryNotReady(err) from err
+
+        self.accounts = await client.async_get_accounts()
+        await self.async_config_entry_first_refresh()
+
+        self.dashboard_url = client.get_dashboard_url()
+        self.user_id = cast(str, client.user_id)

homeassistant/components/ridwell/diagnostics.py

@@ -9,8 +9,8 @@
 from homeassistant.const import CONF_PASSWORD, CONF_UNIQUE_ID, CONF_USERNAME
 from homeassistant.core import HomeAssistant
 
-from . import RidwellData
 from .const import DOMAIN
+from .coordinator import RidwellDataUpdateCoordinator
 
 CONF_TITLE = "title"
 
@@ -27,14 +27,12 @@ async def async_get_config_entry_diagnostics(
     hass: HomeAssistant, entry: ConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    data: RidwellData = hass.data[DOMAIN][entry.entry_id]
+    coordinator: RidwellDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
 
     return async_redact_data(
         {
             "entry": entry.as_dict(),
-            "data": [
-                dataclasses.asdict(event) for event in data.coordinator.data.values()
-            ],
+            "data": [dataclasses.asdict(event) for event in coordinator.data.values()],
         },
         TO_REDACT,
     )

homeassistant/components/ridwell/entity.py

@@ -1,30 +1,36 @@
 """Define a base Ridwell entity."""
 from aioridwell.model import RidwellAccount, RidwellPickupEvent
 
-from homeassistant.helpers.entity import EntityDescription
-from homeassistant.helpers.update_coordinator import (
-    CoordinatorEntity,
-    DataUpdateCoordinator,
-)
+from homeassistant.helpers.device_registry import DeviceEntryType
+from homeassistant.helpers.entity import DeviceInfo, EntityDescription
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
+from .const import DOMAIN
+from .coordinator import RidwellDataUpdateCoordinator
 
-class RidwellEntity(
-    CoordinatorEntity[DataUpdateCoordinator[dict[str, RidwellPickupEvent]]]
-):
+
+class RidwellEntity(CoordinatorEntity[RidwellDataUpdateCoordinator]):
     """Define a base Ridwell entity."""
 
     _attr_has_entity_name = True
 
     def __init__(
         self,
-        coordinator: DataUpdateCoordinator,
+        coordinator: RidwellDataUpdateCoordinator,
         account: RidwellAccount,
         description: EntityDescription,
     ) -> None:
         """Initialize the sensor."""
         super().__init__(coordinator)
 
         self._account = account
+        self._attr_device_info = DeviceInfo(
+            configuration_url=coordinator.dashboard_url,
+            entry_type=DeviceEntryType.SERVICE,
+            identifiers={(DOMAIN, coordinator.user_id)},
+            manufacturer="Ridwell",
+            name="Ridwell",
+        )
         self._attr_unique_id = f"{account.account_id}_{description.key}"
         self.entity_description = description
 

homeassistant/components/ridwell/manifest.json

@@ -3,7 +3,7 @@
   "name": "Ridwell",
   "config_flow": true,
   "documentation": "https://www.home-assistant.io/integrations/ridwell",
-  "requirements": ["aioridwell==2022.11.0"],
+  "requirements": ["aioridwell==2023.01.0"],
   "codeowners": ["@bachya"],
   "iot_class": "cloud_polling",
   "loggers": ["aioridwell"],

homeassistant/components/ridwell/sensor.py

@@ -15,10 +15,9 @@
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from . import RidwellData
 from .const import DOMAIN, SENSOR_TYPE_NEXT_PICKUP
+from .coordinator import RidwellDataUpdateCoordinator
 from .entity import RidwellEntity
 
 ATTR_CATEGORY = "category"
@@ -37,11 +36,11 @@ async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up Ridwell sensors based on a config entry."""
-    data: RidwellData = hass.data[DOMAIN][entry.entry_id]
+    coordinator: RidwellDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
 
     async_add_entities(
-        RidwellSensor(data.coordinator, account, SENSOR_DESCRIPTION)
-        for account in data.accounts.values()
+        RidwellSensor(coordinator, account, SENSOR_DESCRIPTION)
+        for account in coordinator.accounts.values()
     )
 
 
@@ -50,7 +49,7 @@ class RidwellSensor(RidwellEntity, SensorEntity):
 
     def __init__(
         self,
-        coordinator: DataUpdateCoordinator,
+        coordinator: RidwellDataUpdateCoordinator,
         account: RidwellAccount,
         description: SensorEntityDescription,
     ) -> None:

homeassistant/components/ridwell/switch.py

@@ -12,8 +12,8 @@
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import RidwellData
 from .const import DOMAIN
+from .coordinator import RidwellDataUpdateCoordinator
 from .entity import RidwellEntity
 
 SWITCH_TYPE_OPT_IN = "opt_in"
@@ -29,11 +29,11 @@ async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Set up Ridwell sensors based on a config entry."""
-    data: RidwellData = hass.data[DOMAIN][entry.entry_id]
+    coordinator: RidwellDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
 
     async_add_entities(
-        RidwellSwitch(data.coordinator, account, SWITCH_DESCRIPTION)
-        for account in data.accounts.values()
+        RidwellSwitch(coordinator, account, SWITCH_DESCRIPTION)
+        for account in coordinator.accounts.values()
     )
 
 

requirements_all.txt

@@ -258,7 +258,7 @@ aioqsw==0.3.1
 aiorecollect==1.0.8
 
 # homeassistant.components.ridwell
-aioridwell==2022.11.0
+aioridwell==2023.01.0
 
 # homeassistant.components.ruuvi_gateway
 aioruuvigateway==0.0.2

requirements_test_all.txt

@@ -236,7 +236,7 @@ aioqsw==0.3.1
 aiorecollect==1.0.8
 
 # homeassistant.components.ridwell
-aioridwell==2022.11.0
+aioridwell==2023.01.0
 
 # homeassistant.components.ruuvi_gateway
 aioruuvigateway==0.0.2

tests/components/ridwell/conftest.py

@@ -11,8 +11,10 @@
 from tests.common import MockConfigEntry
 
 TEST_ACCOUNT_ID = "12345"
+TEST_DASHBOARD_URL = "https://www.ridwell.com/users/12345/dashboard"
 TEST_PASSWORD = "password"
 TEST_USERNAME = "[email protected]"
+TEST_USER_ID = "12345"
 
 
 @pytest.fixture(name="account")
@@ -44,6 +46,8 @@ def client_fixture(account):
     return Mock(
         async_authenticate=AsyncMock(),
         async_get_accounts=AsyncMock(return_value={TEST_ACCOUNT_ID: account}),
+        get_dashboard_url=Mock(return_value=TEST_DASHBOARD_URL),
+        user_id=TEST_USER_ID,
     )
 
 
@@ -70,7 +74,10 @@ async def mock_aioridwell_fixture(hass, client, config):
     with patch(
         "homeassistant.components.ridwell.config_flow.async_get_client",
         return_value=client,
-    ), patch("homeassistant.components.ridwell.async_get_client", return_value=client):
+    ), patch(
+        "homeassistant.components.ridwell.coordinator.async_get_client",
+        return_value=client,
+    ):
         yield