Change litterrobot integration to cloud_push (#77741)

Co-authored-by: J. Nick Koston <nick@koston.org>

homeassistant/components/litterrobot/__init__.py

@@ -37,7 +37,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Litter-Robot from a config entry."""
     hass.data.setdefault(DOMAIN, {})
     hub = hass.data[DOMAIN][entry.entry_id] = LitterRobotHub(hass, entry.data)
-    await hub.login(load_robots=True)
+    await hub.login(load_robots=True, subscribe_for_updates=True)
 
     if platforms := get_platforms_for_robots(hub.account.robots):
         await hass.config_entries.async_forward_entry_setups(entry, platforms)

homeassistant/components/litterrobot/entity.py

@@ -1,33 +1,24 @@
 """Litter-Robot entities for common data and methods."""
 from __future__ import annotations
 
-from collections.abc import Callable, Coroutine, Iterable
+from collections.abc import Iterable
-from datetime import time
+from typing import Generic, TypeVar
-import logging
-from typing import Any, Generic, TypeVar
 
 from pylitterbot import Robot
-from pylitterbot.exceptions import InvalidCommandException
+from pylitterbot.robot import EVENT_UPDATE
-from typing_extensions import ParamSpec
 
-from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback
+from homeassistant.core import HomeAssistant
-from homeassistant.helpers.entity import DeviceInfo, EntityCategory, EntityDescription
+from homeassistant.helpers.entity import DeviceInfo, EntityDescription
 import homeassistant.helpers.entity_registry as er
-from homeassistant.helpers.event import async_call_later
 from homeassistant.helpers.update_coordinator import (
     CoordinatorEntity,
     DataUpdateCoordinator,
 )
-import homeassistant.util.dt as dt_util
 
 from .const import DOMAIN
 from .hub import LitterRobotHub
 
-_P = ParamSpec("_P")
 _RobotT = TypeVar("_RobotT", bound=Robot)
-_LOGGER = logging.getLogger(__name__)
-
-REFRESH_WAIT_TIME_SECONDS = 8
 
 
 class LitterRobotEntity(
@@ -62,95 +53,10 @@ class LitterRobotEntity(
             sw_version=getattr(self.robot, "firmware", None),
         )
 
-
+    async def async_added_to_hass(self) -> None:
-class LitterRobotControlEntity(LitterRobotEntity[_RobotT]):
+        """Set up a listener for the entity."""
-    """A Litter-Robot entity that can control the unit."""
+        await super().async_added_to_hass()
-
+        self.async_on_remove(self.robot.on(EVENT_UPDATE, self.async_write_ha_state))
-    def __init__(
-        self, robot: _RobotT, hub: LitterRobotHub, description: EntityDescription
-    ) -> None:
-        """Init a Litter-Robot control entity."""
-        super().__init__(robot=robot, hub=hub, description=description)
-        self._refresh_callback: CALLBACK_TYPE | None = None
-
-    async def perform_action_and_refresh(
-        self,
-        action: Callable[_P, Coroutine[Any, Any, bool]],
-        *args: _P.args,
-        **kwargs: _P.kwargs,
-    ) -> bool:
-        """Perform an action and initiates a refresh of the robot data after a few seconds."""
-        success = False
-
-        try:
-            success = await action(*args, **kwargs)
-        except InvalidCommandException as ex:  # pragma: no cover
-            # this exception should only occur if the underlying API for commands changes
-            _LOGGER.error(ex)
-            success = False
-
-        if success:
-            self.async_cancel_refresh_callback()
-            self._refresh_callback = async_call_later(
-                self.hass, REFRESH_WAIT_TIME_SECONDS, self.async_call_later_callback
-            )
-        return success
-
-    async def async_call_later_callback(self, *_: Any) -> None:
-        """Perform refresh request on callback."""
-        self._refresh_callback = None
-        await self.coordinator.async_request_refresh()
-
-    async def async_will_remove_from_hass(self) -> None:
-        """Cancel refresh callback when entity is being removed from hass."""
-        self.async_cancel_refresh_callback()
-
-    @callback
-    def async_cancel_refresh_callback(self) -> None:
-        """Clear the refresh callback if it has not already fired."""
-        if self._refresh_callback is not None:
-            self._refresh_callback()
-            self._refresh_callback = None
-
-    @staticmethod
-    def parse_time_at_default_timezone(time_str: str | None) -> time | None:
-        """Parse a time string and add default timezone."""
-        if time_str is None:
-            return None
-
-        if (parsed_time := dt_util.parse_time(time_str)) is None:  # pragma: no cover
-            return None
-
-        return (
-            dt_util.start_of_local_day()
-            .replace(
-                hour=parsed_time.hour,
-                minute=parsed_time.minute,
-                second=parsed_time.second,
-            )
-            .timetz()
-        )
-
-
-class LitterRobotConfigEntity(LitterRobotControlEntity[_RobotT]):
-    """A Litter-Robot entity that can control configuration of the unit."""
-
-    _attr_entity_category = EntityCategory.CONFIG
-
-    def __init__(
-        self, robot: _RobotT, hub: LitterRobotHub, description: EntityDescription
-    ) -> None:
-        """Init a Litter-Robot control entity."""
-        super().__init__(robot=robot, hub=hub, description=description)
-        self._assumed_state: bool | None = None
-
-    async def perform_action_and_assume_state(
-        self, action: Callable[[bool], Coroutine[Any, Any, bool]], assumed_state: bool
-    ) -> None:
-        """Perform an action and assume the state passed in if call is successful."""
-        if await self.perform_action_and_refresh(action, assumed_state):
-            self._assumed_state = assumed_state
-            self.async_write_ha_state()
 
 
 def async_update_unique_id(

homeassistant/components/litterrobot/hub.py

@@ -19,7 +19,7 @@ from .const import DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 
-UPDATE_INTERVAL_SECONDS = 20
+UPDATE_INTERVAL_SECONDS = 60 * 5
 
 
 class LitterRobotHub:
@@ -43,13 +43,16 @@ class LitterRobotHub:
             update_interval=timedelta(seconds=UPDATE_INTERVAL_SECONDS),
         )
 
-    async def login(self, load_robots: bool = False) -> None:
+    async def login(
+        self, load_robots: bool = False, subscribe_for_updates: bool = False
+    ) -> None:
         """Login to Litter-Robot."""
         try:
             await self.account.connect(
                 username=self._data[CONF_USERNAME],
                 password=self._data[CONF_PASSWORD],
                 load_robots=load_robots,
+                subscribe_for_updates=subscribe_for_updates,
             )
             return
         except LitterRobotLoginException as ex:

homeassistant/components/litterrobot/manifest.json

@@ -6,6 +6,6 @@
   "requirements": ["pylitterbot==2022.9.3"],
   "codeowners": ["@natekspencer", "@tkdrob"],
   "dhcp": [{ "hostname": "litter-robot4" }],
-  "iot_class": "cloud_polling",
+  "iot_class": "cloud_push",
   "loggers": ["pylitterbot"]
 }

homeassistant/components/litterrobot/select.py

@@ -16,10 +16,11 @@ from homeassistant.components.select import (
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import TIME_MINUTES
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .const import DOMAIN
-from .entity import LitterRobotConfigEntity, _RobotT, async_update_unique_id
+from .entity import LitterRobotEntity, _RobotT, async_update_unique_id
 from .hub import LitterRobotHub
 
 _CastTypeT = TypeVar("_CastTypeT", int, float)
@@ -31,10 +32,7 @@ class RequiredKeysMixin(Generic[_RobotT, _CastTypeT]):
 
     current_fn: Callable[[_RobotT], _CastTypeT]
     options_fn: Callable[[_RobotT], list[_CastTypeT]]
-    select_fn: Callable[
+    select_fn: Callable[[_RobotT, str], Coroutine[Any, Any, bool]]
-        [_RobotT, str],
-        tuple[Callable[[_CastTypeT], Coroutine[Any, Any, bool]], _CastTypeT],
-    ]
 
 
 @dataclass
@@ -43,6 +41,8 @@ class RobotSelectEntityDescription(
 ):
     """A class that describes robot select entities."""
 
+    entity_category: EntityCategory = EntityCategory.CONFIG
+
 
 LITTER_ROBOT_SELECT = RobotSelectEntityDescription[LitterRobot, int](
     key="cycle_delay",
@@ -51,7 +51,7 @@ LITTER_ROBOT_SELECT = RobotSelectEntityDescription[LitterRobot, int](
     unit_of_measurement=TIME_MINUTES,
     current_fn=lambda robot: robot.clean_cycle_wait_time_minutes,
     options_fn=lambda robot: robot.VALID_WAIT_TIMES,
-    select_fn=lambda robot, option: (robot.set_wait_time, int(option)),
+    select_fn=lambda robot, option: robot.set_wait_time(int(option)),
 )
 FEEDER_ROBOT_SELECT = RobotSelectEntityDescription[FeederRobot, float](
     key="meal_insert_size",
@@ -60,7 +60,7 @@ FEEDER_ROBOT_SELECT = RobotSelectEntityDescription[FeederRobot, float](
     unit_of_measurement="cups",
     current_fn=lambda robot: robot.meal_insert_size,
     options_fn=lambda robot: robot.VALID_MEAL_INSERT_SIZES,
-    select_fn=lambda robot, option: (robot.set_meal_insert_size, float(option)),
+    select_fn=lambda robot, option: robot.set_meal_insert_size(float(option)),
 )
 
 
@@ -88,7 +88,7 @@ async def async_setup_entry(
 
 
 class LitterRobotSelect(
-    LitterRobotConfigEntity[_RobotT], SelectEntity, Generic[_RobotT, _CastTypeT]
+    LitterRobotEntity[_RobotT], SelectEntity, Generic[_RobotT, _CastTypeT]
 ):
     """Litter-Robot Select."""
 
@@ -112,5 +112,4 @@ class LitterRobotSelect(
 
     async def async_select_option(self, option: str) -> None:
         """Change the selected option."""
-        action, adjusted_option = self.entity_description.select_fn(self.robot, option)
+        await self.entity_description.select_fn(self.robot, option)
-        await self.perform_action_and_refresh(action, adjusted_option)

homeassistant/components/litterrobot/switch.py

@@ -14,10 +14,11 @@ from homeassistant.components.switch import (
 )
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .const import DOMAIN
-from .entity import LitterRobotConfigEntity, _RobotT, async_update_unique_id
+from .entity import LitterRobotEntity, _RobotT, async_update_unique_id
 from .hub import LitterRobotHub
 
 
@@ -26,31 +27,33 @@ class RequiredKeysMixin(Generic[_RobotT]):
     """A class that describes robot switch entity required keys."""
 
     icons: tuple[str, str]
-    set_fn: Callable[[_RobotT], Callable[[bool], Coroutine[Any, Any, bool]]]
+    set_fn: Callable[[_RobotT, bool], Coroutine[Any, Any, bool]]
 
 
 @dataclass
 class RobotSwitchEntityDescription(SwitchEntityDescription, RequiredKeysMixin[_RobotT]):
     """A class that describes robot switch entities."""
 
+    entity_category: EntityCategory = EntityCategory.CONFIG
+
 
 ROBOT_SWITCHES = [
     RobotSwitchEntityDescription[Union[LitterRobot, FeederRobot]](
         key="night_light_mode_enabled",
         name="Night Light Mode",
         icons=("mdi:lightbulb-on", "mdi:lightbulb-off"),
-        set_fn=lambda robot: robot.set_night_light,
+        set_fn=lambda robot, value: robot.set_night_light(value),
     ),
     RobotSwitchEntityDescription[Union[LitterRobot, FeederRobot]](
         key="panel_lock_enabled",
         name="Panel Lockout",
         icons=("mdi:lock", "mdi:lock-open"),
-        set_fn=lambda robot: robot.set_panel_lockout,
+        set_fn=lambda robot, value: robot.set_panel_lockout(value),
     ),
 ]
 
 
-class RobotSwitchEntity(LitterRobotConfigEntity[_RobotT], SwitchEntity):
+class RobotSwitchEntity(LitterRobotEntity[_RobotT], SwitchEntity):
     """Litter-Robot switch entity."""
 
     entity_description: RobotSwitchEntityDescription[_RobotT]
@@ -58,8 +61,6 @@ class RobotSwitchEntity(LitterRobotConfigEntity[_RobotT], SwitchEntity):
     @property
     def is_on(self) -> bool | None:
         """Return true if switch is on."""
-        if self._refresh_callback is not None:
-            return self._assumed_state
         return bool(getattr(self.robot, self.entity_description.key))
 
     @property
@@ -70,13 +71,11 @@ class RobotSwitchEntity(LitterRobotConfigEntity[_RobotT], SwitchEntity):
 
     async def async_turn_on(self, **kwargs: Any) -> None:
         """Turn the switch on."""
-        set_fn = self.entity_description.set_fn
+        await self.entity_description.set_fn(self.robot, True)
-        await self.perform_action_and_assume_state(set_fn(self.robot), True)
 
     async def async_turn_off(self, **kwargs: Any) -> None:
         """Turn the switch off."""
-        set_fn = self.entity_description.set_fn
+        await self.entity_description.set_fn(self.robot, False)
-        await self.perform_action_and_assume_state(set_fn(self.robot), False)
 
 
 async def async_setup_entry(

homeassistant/components/litterrobot/vacuum.py

@@ -1,6 +1,7 @@
 """Support for Litter-Robot "Vacuum"."""
 from __future__ import annotations
 
+from datetime import time
 from typing import Any
 
 from pylitterbot import LitterRobot
@@ -22,9 +23,10 @@ from homeassistant.const import STATE_OFF
 from homeassistant.core import HomeAssistant
 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 .const import DOMAIN
-from .entity import LitterRobotControlEntity, async_update_unique_id
+from .entity import LitterRobotEntity, async_update_unique_id
 from .hub import LitterRobotHub
 
 SERVICE_SET_SLEEP_MODE = "set_sleep_mode"
@@ -70,7 +72,7 @@ async def async_setup_entry(
     )
 
 
-class LitterRobotCleaner(LitterRobotControlEntity[LitterRobot], StateVacuumEntity):
+class LitterRobotCleaner(LitterRobotEntity[LitterRobot], StateVacuumEntity):
     """Litter-Robot "Vacuum" Cleaner."""
 
     _attr_supported_features = (
@@ -95,24 +97,41 @@ class LitterRobotCleaner(LitterRobotControlEntity[LitterRobot], StateVacuumEntit
 
     async def async_turn_on(self, **kwargs: Any) -> None:
         """Turn the cleaner on, starting a clean cycle."""
-        await self.perform_action_and_refresh(self.robot.set_power_status, True)
+        await self.robot.set_power_status(True)
 
     async def async_turn_off(self, **kwargs: Any) -> None:
         """Turn the unit off, stopping any cleaning in progress as is."""
-        await self.perform_action_and_refresh(self.robot.set_power_status, False)
+        await self.robot.set_power_status(False)
 
     async def async_start(self) -> None:
         """Start a clean cycle."""
-        await self.perform_action_and_refresh(self.robot.start_cleaning)
+        await self.robot.start_cleaning()
 
     async def async_set_sleep_mode(
         self, enabled: bool, start_time: str | None = None
     ) -> None:
         """Set the sleep mode."""
-        await self.perform_action_and_refresh(
+        await self.robot.set_sleep_mode(
-            self.robot.set_sleep_mode,
+            enabled, self.parse_time_at_default_timezone(start_time)
-            enabled,
+        )
-            self.parse_time_at_default_timezone(start_time),
+
+    @staticmethod
+    def parse_time_at_default_timezone(time_str: str | None) -> time | None:
+        """Parse a time string and add default timezone."""
+        if time_str is None:
+            return None
+
+        if (parsed_time := dt_util.parse_time(time_str)) is None:  # pragma: no cover
+            return None
+
+        return (
+            dt_util.start_of_local_day()
+            .replace(
+                hour=parsed_time.hour,
+                minute=parsed_time.minute,
+                second=parsed_time.second,
+            )
+            .timetz()
         )
 
     @property

tests/components/litterrobot/test_init.py

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

tests/components/litterrobot/test_select.py

@@ -1,10 +1,7 @@
 """Test the Litter-Robot select entity."""
-from datetime import timedelta
-
 from pylitterbot import LitterRobot3
 import pytest
 
-from homeassistant.components.litterrobot.entity import REFRESH_WAIT_TIME_SECONDS
 from homeassistant.components.select import (
     ATTR_OPTION,
     DOMAIN as PLATFORM_DOMAIN,
@@ -14,12 +11,9 @@ from homeassistant.const import ATTR_ENTITY_ID
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry
 from homeassistant.helpers.entity import EntityCategory
-from homeassistant.util.dt import utcnow
 
 from .conftest import setup_integration
 
-from tests.common import async_fire_time_changed
-
 SELECT_ENTITY_ID = "select.test_clean_cycle_wait_time_minutes"
 
 
@@ -49,8 +43,6 @@ async def test_wait_time_select(hass: HomeAssistant, mock_account):
             blocking=True,
         )
 
-        future = utcnow() + timedelta(seconds=REFRESH_WAIT_TIME_SECONDS)
-        async_fire_time_changed(hass, future)
         assert mock_account.robots[0].set_wait_time.call_count == count
 
 

tests/components/litterrobot/test_switch.py

@@ -1,10 +1,9 @@
 """Test the Litter-Robot switch entity."""
-from datetime import timedelta
 from unittest.mock import MagicMock
 
+from pylitterbot import Robot
 import pytest
 
-from homeassistant.components.litterrobot.entity import REFRESH_WAIT_TIME_SECONDS
 from homeassistant.components.switch import (
     DOMAIN as PLATFORM_DOMAIN,
     SERVICE_TURN_OFF,
@@ -14,12 +13,9 @@ from homeassistant.const import ATTR_ENTITY_ID, STATE_OFF, STATE_ON
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry
 from homeassistant.helpers.entity import EntityCategory
-from homeassistant.util.dt import utcnow
 
 from .conftest import setup_integration
 
-from tests.common import async_fire_time_changed
-
 NIGHT_LIGHT_MODE_ENTITY_ID = "switch.test_night_light_mode"
 PANEL_LOCKOUT_ENTITY_ID = "switch.test_panel_lockout"
 
@@ -39,17 +35,22 @@ async def test_switch(hass: HomeAssistant, mock_account: MagicMock):
 
 
 @pytest.mark.parametrize(
-    "entity_id,robot_command",
+    "entity_id,robot_command,updated_field",
     [
-        (NIGHT_LIGHT_MODE_ENTITY_ID, "set_night_light"),
+        (NIGHT_LIGHT_MODE_ENTITY_ID, "set_night_light", "nightLightActive"),
-        (PANEL_LOCKOUT_ENTITY_ID, "set_panel_lockout"),
+        (PANEL_LOCKOUT_ENTITY_ID, "set_panel_lockout", "panelLockActive"),
     ],
 )
 async def test_on_off_commands(
-    hass: HomeAssistant, mock_account: MagicMock, entity_id: str, robot_command: str
+    hass: HomeAssistant,
+    mock_account: MagicMock,
+    entity_id: str,
+    robot_command: str,
+    updated_field: str,
 ):
     """Test sending commands to the switch."""
     await setup_integration(hass, mock_account, PLATFORM_DOMAIN)
+    robot: Robot = mock_account.robots[0]
 
     state = hass.states.get(entity_id)
     assert state
@@ -57,19 +58,17 @@ async def test_on_off_commands(
     data = {ATTR_ENTITY_ID: entity_id}
 
     count = 0
-    for service in [SERVICE_TURN_ON, SERVICE_TURN_OFF]:
+    for service in (SERVICE_TURN_ON, SERVICE_TURN_OFF):
         count += 1
-
         await hass.services.async_call(
             PLATFORM_DOMAIN,
             service,
             data,
             blocking=True,
         )
+        robot._update_data({updated_field: 1 if service == SERVICE_TURN_ON else 0})
 
-        future = utcnow() + timedelta(seconds=REFRESH_WAIT_TIME_SECONDS)
+        assert getattr(robot, robot_command).call_count == count
-        async_fire_time_changed(hass, future)
-        assert getattr(mock_account.robots[0], robot_command).call_count == count
         state = hass.states.get(entity_id)
         assert state
         assert state.state == STATE_ON if service == SERVICE_TURN_ON else STATE_OFF

tests/components/litterrobot/test_vacuum.py

@@ -1,14 +1,12 @@
 """Test the Litter-Robot vacuum entity."""
 from __future__ import annotations
 
-from datetime import timedelta
 from typing import Any
 from unittest.mock import MagicMock
 
 import pytest
 
 from homeassistant.components.litterrobot import DOMAIN
-from homeassistant.components.litterrobot.entity import REFRESH_WAIT_TIME_SECONDS
 from homeassistant.components.litterrobot.vacuum import SERVICE_SET_SLEEP_MODE
 from homeassistant.components.vacuum import (
     ATTR_STATUS,
@@ -22,13 +20,10 @@ from homeassistant.components.vacuum import (
 from homeassistant.const import ATTR_ENTITY_ID
 from homeassistant.core import HomeAssistant
 import homeassistant.helpers.entity_registry as er
-from homeassistant.util.dt import utcnow
 
 from .common import VACUUM_ENTITY_ID
 from .conftest import setup_integration
 
-from tests.common import async_fire_time_changed
-
 VACUUM_UNIQUE_ID_OLD = "LR3C012345-Litter Box"
 VACUUM_UNIQUE_ID_NEW = "LR3C012345-litter_box"
 
@@ -141,7 +136,5 @@ async def test_commands(
         data,
         blocking=True,
     )
-    future = utcnow() + timedelta(seconds=REFRESH_WAIT_TIME_SECONDS)
-    async_fire_time_changed(hass, future)
     getattr(mock_account.robots[0], command).assert_called_once()
     assert (f"'{DOMAIN}.{service}' service is deprecated" in caplog.text) is deprecated