Further generalize base Ridwell entity (#85486)

.coveragerc

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

homeassistant/components/ridwell/__init__.py

@@ -15,12 +15,7 @@ from homeassistant.const import CONF_PASSWORD, CONF_USERNAME, 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.entity import EntityDescription
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
-from homeassistant.helpers.update_coordinator import (
-    CoordinatorEntity,
-    DataUpdateCoordinator,
-    UpdateFailed,
-)
 
 from .const import DOMAIN, LOGGER, SENSOR_TYPE_NEXT_PICKUP
 
@@ -34,7 +29,7 @@ class RidwellData:
     """Define an object to be stored in `hass.data`."""
 
     accounts: dict[str, RidwellAccount]
-    coordinator: DataUpdateCoordinator
+    coordinator: DataUpdateCoordinator[dict[str, RidwellPickupEvent]]
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
@@ -70,7 +65,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
         return data
 
-    coordinator = DataUpdateCoordinator(
+    coordinator: DataUpdateCoordinator[
+        dict[str, RidwellPickupEvent]
+    ] = DataUpdateCoordinator(
         hass,
         LOGGER,
         name=entry.title,
@@ -79,8 +76,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     )
 
     await coordinator.async_config_entry_first_refresh()
-    hass.data.setdefault(DOMAIN, {})
+    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = RidwellData(
-    hass.data[DOMAIN][entry.entry_id] = RidwellData(
         accounts=accounts, coordinator=coordinator
     )
 
@@ -91,8 +87,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
+    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-    if unload_ok:
         hass.data[DOMAIN].pop(entry.entry_id)
 
     return unload_ok
@@ -121,22 +116,3 @@ async def async_migrate_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     LOGGER.info("Migration to version %s successful", version)
 
     return True
-
-
-class RidwellEntity(CoordinatorEntity):
-    """Define a base Ridwell entity."""
-
-    _attr_has_entity_name = True
-
-    def __init__(
-        self,
-        coordinator: DataUpdateCoordinator,
-        account: RidwellAccount,
-        description: EntityDescription,
-    ) -> None:
-        """Initialize the sensor."""
-        super().__init__(coordinator)
-
-        self._account = account
-        self._attr_unique_id = f"{account.account_id}_{description.key}"
-        self.entity_description = description

homeassistant/components/ridwell/entity.py

@@ -0,0 +1,34 @@
+"""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,
+)
+
+
+class RidwellEntity(
+    CoordinatorEntity[DataUpdateCoordinator[dict[str, RidwellPickupEvent]]]
+):
+    """Define a base Ridwell entity."""
+
+    _attr_has_entity_name = True
+
+    def __init__(
+        self,
+        coordinator: DataUpdateCoordinator,
+        account: RidwellAccount,
+        description: EntityDescription,
+    ) -> None:
+        """Initialize the sensor."""
+        super().__init__(coordinator)
+
+        self._account = account
+        self._attr_unique_id = f"{account.account_id}_{description.key}"
+        self.entity_description = description
+
+    @property
+    def next_pickup_event(self) -> RidwellPickupEvent:
+        """Get the next pickup event."""
+        return self.coordinator.data[self._account.account_id]

homeassistant/components/ridwell/sensor.py

@@ -2,10 +2,10 @@
 from __future__ import annotations
 
 from collections.abc import Mapping
-from datetime import date, datetime
+from datetime import date
 from typing import Any
 
-from aioridwell.model import RidwellAccount, RidwellPickupEvent
+from aioridwell.model import RidwellAccount
 
 from homeassistant.components.sensor import (
     SensorDeviceClass,
@@ -15,11 +15,11 @@ from homeassistant.components.sensor import (
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
-from homeassistant.helpers.typing import StateType
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from . import RidwellData, RidwellEntity
+from . import RidwellData
 from .const import DOMAIN, SENSOR_TYPE_NEXT_PICKUP
+from .entity import RidwellEntity
 
 ATTR_CATEGORY = "category"
 ATTR_PICKUP_STATE = "pickup_state"
@@ -40,10 +40,8 @@ async def async_setup_entry(
     data: RidwellData = hass.data[DOMAIN][entry.entry_id]
 
     async_add_entities(
-        [
         RidwellSensor(data.coordinator, account, SENSOR_DESCRIPTION)
         for account in data.accounts.values()
-        ]
     )
 
 
@@ -64,14 +62,12 @@ class RidwellSensor(RidwellEntity, SensorEntity):
     @property
     def extra_state_attributes(self) -> Mapping[str, Any]:
         """Return entity specific state attributes."""
-        event = self.coordinator.data[self._account.account_id]
-
         attrs: dict[str, Any] = {
             ATTR_PICKUP_TYPES: {},
-            ATTR_PICKUP_STATE: event.state.value,
+            ATTR_PICKUP_STATE: self.next_pickup_event.state.value,
         }
 
-        for pickup in event.pickups:
+        for pickup in self.next_pickup_event.pickups:
             if pickup.name not in attrs[ATTR_PICKUP_TYPES]:
                 attrs[ATTR_PICKUP_TYPES][pickup.name] = {
                     ATTR_CATEGORY: pickup.category.value,
@@ -86,7 +82,6 @@ class RidwellSensor(RidwellEntity, SensorEntity):
         return attrs
 
     @property
-    def native_value(self) -> StateType | date | datetime:
+    def native_value(self) -> date:
         """Return the value reported by the sensor."""
-        event: RidwellPickupEvent = self.coordinator.data[self._account.account_id]
+        return self.next_pickup_event.pickup_date
-        return event.pickup_date

homeassistant/components/ridwell/switch.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 from typing import Any
 
 from aioridwell.errors import RidwellError
-from aioridwell.model import EventState, RidwellPickupEvent
+from aioridwell.model import EventState
 
 from homeassistant.components.switch import SwitchEntity, SwitchEntityDescription
 from homeassistant.config_entries import ConfigEntry
@@ -12,14 +12,15 @@ from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import RidwellData, RidwellEntity
+from . import RidwellData
 from .const import DOMAIN
+from .entity import RidwellEntity
 
 SWITCH_TYPE_OPT_IN = "opt_in"
 
 SWITCH_DESCRIPTION = SwitchEntityDescription(
     key=SWITCH_TYPE_OPT_IN,
-    name="Opt-In to Next Pickup",
+    name="Opt-in to next pickup",
     icon="mdi:calendar-check",
 )
 
@@ -31,10 +32,8 @@ async def async_setup_entry(
     data: RidwellData = hass.data[DOMAIN][entry.entry_id]
 
     async_add_entities(
-        [
         RidwellSwitch(data.coordinator, account, SWITCH_DESCRIPTION)
         for account in data.accounts.values()
-        ]
     )
 
 
@@ -44,15 +43,12 @@ class RidwellSwitch(RidwellEntity, SwitchEntity):
     @property
     def is_on(self) -> bool:
         """Return True if entity is on."""
-        event: RidwellPickupEvent = self.coordinator.data[self._account.account_id]
+        return self.next_pickup_event.state == EventState.SCHEDULED
-        return event.state == EventState.SCHEDULED
 
     async def async_turn_off(self, **kwargs: Any) -> None:
         """Turn the switch off."""
-        event: RidwellPickupEvent = self.coordinator.data[self._account.account_id]
-
         try:
-            await event.async_opt_out()
+            await self.next_pickup_event.async_opt_out()
         except RidwellError as err:
             raise HomeAssistantError(f"Error while opting out: {err}") from err
 
@@ -60,10 +56,8 @@ class RidwellSwitch(RidwellEntity, SwitchEntity):
 
     async def async_turn_on(self, **kwargs: Any) -> None:
         """Turn the switch on."""
-        event: RidwellPickupEvent = self.coordinator.data[self._account.account_id]
-
         try:
-            await event.async_opt_in()
+            await self.next_pickup_event.async_opt_in()
         except RidwellError as err:
             raise HomeAssistantError(f"Error while opting in: {err}") from err