Convert LitterRobotHub to a DataUpdateCoordinator (#136283)

homeassistant/components/litterrobot/__init__.py

@@ -4,15 +4,12 @@ from __future__ import annotations
 
 from pylitterbot import FeederRobot, LitterRobot, LitterRobot3, LitterRobot4, Robot
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.device_registry import DeviceEntry
 
 from .const import DOMAIN
-from .hub import LitterRobotHub
+from .coordinator import LitterRobotConfigEntry, LitterRobotDataUpdateCoordinator
-
-type LitterRobotConfigEntry = ConfigEntry[LitterRobotHub]
 
 PLATFORMS_BY_TYPE = {
     Robot: (
@@ -41,11 +38,11 @@ def get_platforms_for_robots(robots: list[Robot]) -> set[Platform]:
 
 async def async_setup_entry(hass: HomeAssistant, entry: LitterRobotConfigEntry) -> bool:
     """Set up Litter-Robot from a config entry."""
-    hub = LitterRobotHub(hass, entry.data)
+    coordinator = LitterRobotDataUpdateCoordinator(hass, entry)
-    await hub.login(load_robots=True, subscribe_for_updates=True)
+    await coordinator.async_config_entry_first_refresh()
-    entry.runtime_data = hub
+    entry.runtime_data = coordinator
 
-    if platforms := get_platforms_for_robots(hub.account.robots):
+    if platforms := get_platforms_for_robots(coordinator.account.robots):
         await hass.config_entries.async_forward_entry_setups(entry, platforms)
     return True
 

homeassistant/components/litterrobot/binary_sensor.py

@@ -17,7 +17,7 @@ from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import LitterRobotConfigEntry
+from .coordinator import LitterRobotConfigEntry
 from .entity import LitterRobotEntity, _RobotT
 
 
@@ -66,10 +66,12 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Litter-Robot binary sensors using config entry."""
-    hub = entry.runtime_data
+    coordinator = entry.runtime_data
     async_add_entities(
-        LitterRobotBinarySensorEntity(robot=robot, hub=hub, description=description)
+        LitterRobotBinarySensorEntity(
-        for robot in hub.account.robots
+            robot=robot, coordinator=coordinator, description=description
+        )
+        for robot in coordinator.account.robots
         for robot_type, entity_descriptions in BINARY_SENSOR_MAP.items()
         if isinstance(robot, robot_type)
         for description in entity_descriptions

homeassistant/components/litterrobot/button.py

@@ -13,7 +13,7 @@ from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import LitterRobotConfigEntry
+from .coordinator import LitterRobotConfigEntry
 from .entity import LitterRobotEntity, _RobotT
 
 
@@ -51,14 +51,15 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Litter-Robot cleaner using config entry."""
-    hub = entry.runtime_data
+    coordinator = entry.runtime_data
-    entities = [
+    async_add_entities(
-        LitterRobotButtonEntity(robot=robot, hub=hub, description=description)
+        LitterRobotButtonEntity(
-        for robot in hub.account.robots
+            robot=robot, coordinator=coordinator, description=description
+        )
+        for robot in coordinator.account.robots
         for robot_type, description in ROBOT_BUTTON_MAP.items()
         if isinstance(robot, robot_type)
-    ]
+    )
-    async_add_entities(entities)
 
 
 class LitterRobotButtonEntity(LitterRobotEntity[_RobotT], ButtonEntity):
@@ -69,4 +70,4 @@ class LitterRobotButtonEntity(LitterRobotEntity[_RobotT], ButtonEntity):
     async def async_press(self) -> None:
         """Press the button."""
         await self.entity_description.press_fn(self.robot)
-        self.coordinator.async_set_updated_data(True)
+        self.coordinator.async_set_updated_data(None)

homeassistant/components/litterrobot/coordinator.py

@@ -1,64 +1,66 @@
-"""A wrapper 'hub' for the Litter-Robot API."""
+"""The Litter-Robot coordinator."""
 
 from __future__ import annotations
 
-from collections.abc import Generator, Mapping
+from collections.abc import Generator
 from datetime import timedelta
 import logging
-from typing import Any
 
 from pylitterbot import Account, FeederRobot, LitterRobot
 from pylitterbot.exceptions import LitterRobotException, LitterRobotLoginException
 
+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.exceptions import ConfigEntryAuthFailed
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
 from .const import DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 
-UPDATE_INTERVAL_SECONDS = 60 * 5
+UPDATE_INTERVAL = timedelta(minutes=5)
+
+type LitterRobotConfigEntry = ConfigEntry[LitterRobotDataUpdateCoordinator]
 
 
-class LitterRobotHub:
+class LitterRobotDataUpdateCoordinator(DataUpdateCoordinator[None]):
-    """A Litter-Robot hub wrapper class."""
+    """The Litter-Robot data update coordinator."""
 
-    def __init__(self, hass: HomeAssistant, data: Mapping[str, Any]) -> None:
+    config_entry: LitterRobotConfigEntry
-        """Initialize the Litter-Robot hub."""
-        self._data = data
-        self.account = Account(websession=async_get_clientsession(hass))
 
-        async def _async_update_data() -> bool:
+    def __init__(
-            """Update all device states from the Litter-Robot API."""
+        self, hass: HomeAssistant, config_entry: LitterRobotConfigEntry
-            await self.account.refresh_robots()
+    ) -> None:
-            return True
+        """Initialize the Litter-Robot data update coordinator."""
-
+        super().__init__(
-        self.coordinator = DataUpdateCoordinator(
             hass,
             _LOGGER,
+            config_entry=config_entry,
             name=DOMAIN,
-            update_method=_async_update_data,
+            update_interval=UPDATE_INTERVAL,
-            update_interval=timedelta(seconds=UPDATE_INTERVAL_SECONDS),
         )
 
-    async def login(
+        self.account = Account(websession=async_get_clientsession(hass))
-        self, load_robots: bool = False, subscribe_for_updates: bool = False
+
-    ) -> None:
+    async def _async_update_data(self) -> None:
-        """Login to Litter-Robot."""
+        """Update all device states from the Litter-Robot API."""
+        await self.account.refresh_robots()
+
+    async def _async_setup(self) -> None:
+        """Set up the coordinator."""
         try:
             await self.account.connect(
-                username=self._data[CONF_USERNAME],
+                username=self.config_entry.data[CONF_USERNAME],
-                password=self._data[CONF_PASSWORD],
+                password=self.config_entry.data[CONF_PASSWORD],
-                load_robots=load_robots,
+                load_robots=True,
-                subscribe_for_updates=subscribe_for_updates,
+                subscribe_for_updates=True,
             )
         except LitterRobotLoginException as ex:
             raise ConfigEntryAuthFailed("Invalid credentials") from ex
         except LitterRobotException as ex:
-            raise ConfigEntryNotReady("Unable to connect to Litter-Robot API") from ex
+            raise UpdateFailed("Unable to connect to Litter-Robot API") from ex
 
     def litter_robots(self) -> Generator[LitterRobot]:
         """Get Litter-Robots from the account."""

homeassistant/components/litterrobot/entity.py

@@ -9,44 +9,39 @@ from pylitterbot.robot import EVENT_UPDATE
 
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity import EntityDescription
-from homeassistant.helpers.update_coordinator import (
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
-    CoordinatorEntity,
-    DataUpdateCoordinator,
-)
 
 from .const import DOMAIN
-from .hub import LitterRobotHub
+from .coordinator import LitterRobotDataUpdateCoordinator
 
 _RobotT = TypeVar("_RobotT", bound=Robot)
 
 
 class LitterRobotEntity(
-    CoordinatorEntity[DataUpdateCoordinator[bool]], Generic[_RobotT]
+    CoordinatorEntity[LitterRobotDataUpdateCoordinator], Generic[_RobotT]
 ):
     """Generic Litter-Robot entity representing common data and methods."""
 
     _attr_has_entity_name = True
 
     def __init__(
-        self, robot: _RobotT, hub: LitterRobotHub, description: EntityDescription
+        self,
+        robot: _RobotT,
+        coordinator: LitterRobotDataUpdateCoordinator,
+        description: EntityDescription,
     ) -> None:
         """Pass coordinator to CoordinatorEntity."""
-        super().__init__(hub.coordinator)
+        super().__init__(coordinator)
         self.robot = robot
-        self.hub = hub
         self.entity_description = description
-        self._attr_unique_id = f"{self.robot.serial}-{description.key}"
+        self._attr_unique_id = f"{robot.serial}-{description.key}"
-
+        self._attr_device_info = DeviceInfo(
-    @property
+            identifiers={(DOMAIN, robot.serial)},
-    def device_info(self) -> DeviceInfo:
+            manufacturer="Whisker",
-        """Return the device information for a Litter-Robot."""
+            model=robot.model,
-        assert self.robot.serial
+            name=robot.name,
-        return DeviceInfo(
+            serial_number=robot.serial,
-            identifiers={(DOMAIN, self.robot.serial)},
+            sw_version=getattr(robot, "firmware", None),
-            manufacturer="Litter-Robot",
-            model=self.robot.model,
-            name=self.robot.name,
-            sw_version=getattr(self.robot, "firmware", None),
         )
 
     async def async_added_to_hass(self) -> None:

homeassistant/components/litterrobot/select.py

@@ -14,9 +14,8 @@ from homeassistant.const import EntityCategory, UnitOfTime
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import LitterRobotConfigEntry
+from .coordinator import LitterRobotConfigEntry, LitterRobotDataUpdateCoordinator
 from .entity import LitterRobotEntity, _RobotT
-from .hub import LitterRobotHub
 
 _CastTypeT = TypeVar("_CastTypeT", int, float, str)
 
@@ -72,14 +71,15 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Litter-Robot selects using config entry."""
-    hub = entry.runtime_data
+    coordinator = entry.runtime_data
-    entities = [
+    async_add_entities(
-        LitterRobotSelectEntity(robot=robot, hub=hub, description=description)
+        LitterRobotSelectEntity(
-        for robot in hub.account.robots
+            robot=robot, coordinator=coordinator, description=description
+        )
+        for robot in coordinator.account.robots
         for robot_type, description in ROBOT_SELECT_MAP.items()
         if isinstance(robot, robot_type)
-    ]
+    )
-    async_add_entities(entities)
 
 
 class LitterRobotSelectEntity(
@@ -92,11 +92,11 @@ class LitterRobotSelectEntity(
     def __init__(
         self,
         robot: _RobotT,
-        hub: LitterRobotHub,
+        coordinator: LitterRobotDataUpdateCoordinator,
         description: RobotSelectEntityDescription[_RobotT, _CastTypeT],
     ) -> None:
         """Initialize a Litter-Robot select entity."""
-        super().__init__(robot, hub, description)
+        super().__init__(robot, coordinator, description)
         options = self.entity_description.options_fn(self.robot)
         self._attr_options = list(map(str, options))
 

homeassistant/components/litterrobot/sensor.py

@@ -19,7 +19,7 @@ from homeassistant.const import PERCENTAGE, EntityCategory, UnitOfMass
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import LitterRobotConfigEntry
+from .coordinator import LitterRobotConfigEntry
 from .entity import LitterRobotEntity, _RobotT
 
 
@@ -159,12 +159,13 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Litter-Robot sensors using config entry."""
-    hub = entry.runtime_data
+    coordinator = entry.runtime_data
-    entities = [
+    async_add_entities(
-        LitterRobotSensorEntity(robot=robot, hub=hub, description=description)
+        LitterRobotSensorEntity(
-        for robot in hub.account.robots
+            robot=robot, coordinator=coordinator, description=description
+        )
+        for robot in coordinator.account.robots
         for robot_type, entity_descriptions in ROBOT_SENSOR_MAP.items()
         if isinstance(robot, robot_type)
         for description in entity_descriptions
-    ]
+    )
-    async_add_entities(entities)

homeassistant/components/litterrobot/switch.py

@@ -13,7 +13,7 @@ from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import LitterRobotConfigEntry
+from .coordinator import LitterRobotConfigEntry
 from .entity import LitterRobotEntity, _RobotT
 
 
@@ -48,14 +48,13 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Litter-Robot switches using config entry."""
-    hub = entry.runtime_data
+    coordinator = entry.runtime_data
-    entities = [
+    async_add_entities(
-        RobotSwitchEntity(robot=robot, hub=hub, description=description)
+        RobotSwitchEntity(robot=robot, coordinator=coordinator, description=description)
         for description in ROBOT_SWITCHES
-        for robot in hub.account.robots
+        for robot in coordinator.account.robots
         if isinstance(robot, (LitterRobot, FeederRobot))
-    ]
+    )
-    async_add_entities(entities)
 
 
 class RobotSwitchEntity(LitterRobotEntity[_RobotT], SwitchEntity):

homeassistant/components/litterrobot/time.py

@@ -15,7 +15,7 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 import homeassistant.util.dt as dt_util
 
-from . import LitterRobotConfigEntry
+from .coordinator import LitterRobotConfigEntry
 from .entity import LitterRobotEntity, _RobotT
 
 
@@ -52,15 +52,15 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Litter-Robot cleaner using config entry."""
-    hub = entry.runtime_data
+    coordinator = entry.runtime_data
     async_add_entities(
-        [
         LitterRobotTimeEntity(
-                robot=robot, hub=hub, description=LITTER_ROBOT_3_SLEEP_START
+            robot=robot,
+            coordinator=coordinator,
+            description=LITTER_ROBOT_3_SLEEP_START,
         )
-            for robot in hub.litter_robots()
+        for robot in coordinator.litter_robots()
         if isinstance(robot, LitterRobot3)
-        ]
     )
 
 

homeassistant/components/litterrobot/update.py

@@ -17,7 +17,7 @@ from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import LitterRobotConfigEntry
+from .coordinator import LitterRobotConfigEntry
 from .entity import LitterRobotEntity
 
 SCAN_INTERVAL = timedelta(days=1)
@@ -34,12 +34,14 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Litter-Robot update platform."""
-    hub = entry.runtime_data
+    coordinator = entry.runtime_data
-    entities = [
+    entities = (
-        RobotUpdateEntity(robot=robot, hub=hub, description=FIRMWARE_UPDATE_ENTITY)
+        RobotUpdateEntity(
-        for robot in hub.litter_robots()
+            robot=robot, coordinator=coordinator, description=FIRMWARE_UPDATE_ENTITY
+        )
+        for robot in coordinator.litter_robots()
         if isinstance(robot, LitterRobot4)
-    ]
+    )
     async_add_entities(entities, True)
 
 

homeassistant/components/litterrobot/vacuum.py

@@ -20,7 +20,7 @@ from homeassistant.helpers import config_validation as cv, entity_platform
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 import homeassistant.util.dt as dt_util
 
-from . import LitterRobotConfigEntry
+from .coordinator import LitterRobotConfigEntry
 from .entity import LitterRobotEntity
 
 SERVICE_SET_SLEEP_MODE = "set_sleep_mode"
@@ -49,12 +49,13 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Litter-Robot cleaner using config entry."""
-    hub = entry.runtime_data
+    coordinator = entry.runtime_data
-    entities = [
+    async_add_entities(
-        LitterRobotCleaner(robot=robot, hub=hub, description=LITTER_BOX_ENTITY)
+        LitterRobotCleaner(
-        for robot in hub.litter_robots()
+            robot=robot, coordinator=coordinator, description=LITTER_BOX_ENTITY
-    ]
+        )
-    async_add_entities(entities)
+        for robot in coordinator.litter_robots()
+    )
 
     platform = entity_platform.async_get_current_platform()
     platform.async_register_entity_service(

tests/components/litterrobot/conftest.py

@@ -117,7 +117,7 @@ def mock_account_with_side_effects() -> MagicMock:
 async def setup_integration(
     hass: HomeAssistant, mock_account: MagicMock, platform_domain: str | None = None
 ) -> MockConfigEntry:
-    """Load a Litter-Robot platform with the provided hub."""
+    """Load a Litter-Robot platform with the provided coordinator."""
     entry = MockConfigEntry(
         domain=litterrobot.DOMAIN,
         data=CONFIG[litterrobot.DOMAIN],
@@ -126,7 +126,7 @@ async def setup_integration(
 
     with (
         patch(
-            "homeassistant.components.litterrobot.hub.Account",
+            "homeassistant.components.litterrobot.coordinator.Account",
             return_value=mock_account,
         ),
         patch(

tests/components/litterrobot/test_init.py

@@ -63,7 +63,7 @@ async def test_entry_not_setup(
     entry.add_to_hass(hass)
 
     with patch(
-        "homeassistant.components.litterrobot.hub.Account.connect",
+        "homeassistant.components.litterrobot.coordinator.Account.connect",
         side_effect=side_effect,
     ):
         await hass.config_entries.async_setup(entry.entry_id)