Add set_wait_time command support to Litter-Robot (#48300)

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

homeassistant/components/litterrobot/__init__.py

@@ -30,10 +30,11 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry):
     except LitterRobotException as ex:
         raise ConfigEntryNotReady from ex
 
-    for platform in PLATFORMS:
-        hass.async_create_task(
-            hass.config_entries.async_forward_entry_setup(entry, platform)
-        )
+    if hub.account.robots:
+        for platform in PLATFORMS:
+            hass.async_create_task(
+                hass.config_entries.async_forward_entry_setup(entry, platform)
+            )
 
     return True
 

homeassistant/components/litterrobot/entity.py

@@ -0,0 +1,113 @@
+"""Litter-Robot entities for common data and methods."""
+from __future__ import annotations
+
+from datetime import time
+import logging
+from types import MethodType
+from typing import Any
+
+from pylitterbot import Robot
+from pylitterbot.exceptions import InvalidCommandException
+
+from homeassistant.core import callback
+from homeassistant.helpers.event import async_call_later
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
+import homeassistant.util.dt as dt_util
+
+from .const import DOMAIN
+from .hub import LitterRobotHub
+
+_LOGGER = logging.getLogger(__name__)
+
+REFRESH_WAIT_TIME_SECONDS = 8
+
+
+class LitterRobotEntity(CoordinatorEntity):
+    """Generic Litter-Robot entity representing common data and methods."""
+
+    def __init__(self, robot: Robot, entity_type: str, hub: LitterRobotHub) -> None:
+        """Pass coordinator to CoordinatorEntity."""
+        super().__init__(hub.coordinator)
+        self.robot = robot
+        self.entity_type = entity_type
+        self.hub = hub
+
+    @property
+    def name(self) -> str:
+        """Return the name of this entity."""
+        return f"{self.robot.name} {self.entity_type}"
+
+    @property
+    def unique_id(self) -> str:
+        """Return a unique ID."""
+        return f"{self.robot.serial}-{self.entity_type}"
+
+    @property
+    def device_info(self) -> dict[str, Any]:
+        """Return the device information for a Litter-Robot."""
+        return {
+            "identifiers": {(DOMAIN, self.robot.serial)},
+            "name": self.robot.name,
+            "manufacturer": "Litter-Robot",
+            "model": self.robot.model,
+        }
+
+
+class LitterRobotControlEntity(LitterRobotEntity):
+    """A Litter-Robot entity that can control the unit."""
+
+    def __init__(self, robot: Robot, entity_type: str, hub: LitterRobotHub) -> None:
+        """Init a Litter-Robot control entity."""
+        super().__init__(robot=robot, entity_type=entity_type, hub=hub)
+        self._refresh_callback = None
+
+    async def perform_action_and_refresh(
+        self, action: MethodType, *args: Any, **kwargs: Any
+    ) -> bool:
+        """Perform an action and initiates a refresh of the robot data after a few seconds."""
+
+        try:
+            await action(*args, **kwargs)
+        except InvalidCommandException as ex:
+            _LOGGER.error(ex)
+            return False
+
+        self.async_cancel_refresh_callback()
+        self._refresh_callback = async_call_later(
+            self.hass, REFRESH_WAIT_TIME_SECONDS, self.async_call_later_callback
+        )
+        return True
+
+    async def async_call_later_callback(self, *_) -> 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):
+        """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) -> time | None:
+        """Parse a time string and add default timezone."""
+        parsed_time = dt_util.parse_time(time_str)
+
+        if parsed_time is None:
+            return None
+
+        return (
+            dt_util.start_of_local_day()
+            .replace(
+                hour=parsed_time.hour,
+                minute=parsed_time.minute,
+                second=parsed_time.second,
+            )
+            .timetz()
+        )

homeassistant/components/litterrobot/hub.py

@@ -1,41 +1,31 @@
-"""A wrapper 'hub' for the Litter-Robot API and base entity for common attributes."""
-from __future__ import annotations
-
-from datetime import time, timedelta
+"""A wrapper 'hub' for the Litter-Robot API."""
+from datetime import timedelta
 import logging
-from types import MethodType
-from typing import Any
 
-import pylitterbot
+from pylitterbot import Account
 from pylitterbot.exceptions import LitterRobotException, LitterRobotLoginException
 
 from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.event import async_call_later
-from homeassistant.helpers.update_coordinator import (
-    CoordinatorEntity,
-    DataUpdateCoordinator,
-)
-import homeassistant.util.dt as dt_util
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
 from .const import DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 
-REFRESH_WAIT_TIME = 12
-UPDATE_INTERVAL = 10
+UPDATE_INTERVAL_SECONDS = 10
 
 
 class LitterRobotHub:
     """A Litter-Robot hub wrapper class."""
 
-    def __init__(self, hass: HomeAssistant, data: dict):
+    def __init__(self, hass: HomeAssistant, data: dict) -> None:
         """Initialize the Litter-Robot hub."""
         self._data = data
         self.account = None
         self.logged_in = False
 
-        async def _async_update_data():
+        async def _async_update_data() -> bool:
             """Update all device states from the Litter-Robot API."""
             await self.account.refresh_robots()
             return True
@@ -45,13 +35,13 @@ class LitterRobotHub:
             _LOGGER,
             name=DOMAIN,
             update_method=_async_update_data,
-            update_interval=timedelta(seconds=UPDATE_INTERVAL),
+            update_interval=timedelta(seconds=UPDATE_INTERVAL_SECONDS),
         )
 
-    async def login(self, load_robots: bool = False):
+    async def login(self, load_robots: bool = False) -> None:
         """Login to Litter-Robot."""
         self.logged_in = False
-        self.account = pylitterbot.Account()
+        self.account = Account()
         try:
             await self.account.connect(
                 username=self._data[CONF_USERNAME],
@@ -66,61 +56,3 @@ class LitterRobotHub:
         except LitterRobotException as ex:
             _LOGGER.error("Unable to connect to Litter-Robot API")
             raise ex
-
-
-class LitterRobotEntity(CoordinatorEntity):
-    """Generic Litter-Robot entity representing common data and methods."""
-
-    def __init__(self, robot: pylitterbot.Robot, entity_type: str, hub: LitterRobotHub):
-        """Pass coordinator to CoordinatorEntity."""
-        super().__init__(hub.coordinator)
-        self.robot = robot
-        self.entity_type = entity_type
-        self.hub = hub
-
-    @property
-    def name(self):
-        """Return the name of this entity."""
-        return f"{self.robot.name} {self.entity_type}"
-
-    @property
-    def unique_id(self):
-        """Return a unique ID."""
-        return f"{self.robot.serial}-{self.entity_type}"
-
-    @property
-    def device_info(self):
-        """Return the device information for a Litter-Robot."""
-        return {
-            "identifiers": {(DOMAIN, self.robot.serial)},
-            "name": self.robot.name,
-            "manufacturer": "Litter-Robot",
-            "model": self.robot.model,
-        }
-
-    async def perform_action_and_refresh(self, action: MethodType, *args: Any):
-        """Perform an action and initiates a refresh of the robot data after a few seconds."""
-
-        async def async_call_later_callback(*_) -> None:
-            await self.hub.coordinator.async_request_refresh()
-
-        await action(*args)
-        async_call_later(self.hass, REFRESH_WAIT_TIME, async_call_later_callback)
-
-    @staticmethod
-    def parse_time_at_default_timezone(time_str: str) -> time | None:
-        """Parse a time string and add default timezone."""
-        parsed_time = dt_util.parse_time(time_str)
-
-        if parsed_time is None:
-            return None
-
-        return (
-            dt_util.start_of_local_day()
-            .replace(
-                hour=parsed_time.hour,
-                minute=parsed_time.minute,
-                second=parsed_time.second,
-            )
-            .timetz()
-        )

homeassistant/components/litterrobot/manifest.json

@@ -3,6 +3,6 @@
   "name": "Litter-Robot",
   "config_flow": true,
   "documentation": "https://www.home-assistant.io/integrations/litterrobot",
-  "requirements": ["pylitterbot==2021.2.8"],
+  "requirements": ["pylitterbot==2021.3.1"],
   "codeowners": ["@natekspencer"]
 }

homeassistant/components/litterrobot/sensor.py

@@ -1,13 +1,19 @@
 """Support for Litter-Robot sensors."""
 from __future__ import annotations
 
+from typing import Callable
+
 from pylitterbot.robot import Robot
 
 from homeassistant.components.sensor import SensorEntity
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import DEVICE_CLASS_TIMESTAMP, PERCENTAGE
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity import Entity
 
 from .const import DOMAIN
-from .hub import LitterRobotEntity, LitterRobotHub
+from .entity import LitterRobotEntity
+from .hub import LitterRobotHub
 
 
 def icon_for_gauge_level(gauge_level: int | None = None, offset: int = 0) -> str:
@@ -22,66 +28,76 @@ def icon_for_gauge_level(gauge_level: int | None = None, offset: int = 0) -> str
 
 
 class LitterRobotPropertySensor(LitterRobotEntity, SensorEntity):
-    """Litter-Robot property sensors."""
+    """Litter-Robot property sensor."""
 
     def __init__(
         self, robot: Robot, entity_type: str, hub: LitterRobotHub, sensor_attribute: str
-    ):
-        """Pass coordinator to CoordinatorEntity."""
+    ) -> None:
+        """Pass robot, entity_type and hub to LitterRobotEntity."""
         super().__init__(robot, entity_type, hub)
         self.sensor_attribute = sensor_attribute
 
     @property
-    def state(self):
+    def state(self) -> str:
         """Return the state."""
         return getattr(self.robot, self.sensor_attribute)
 
 
 class LitterRobotWasteSensor(LitterRobotPropertySensor):
-    """Litter-Robot sensors."""
+    """Litter-Robot waste sensor."""
 
     @property
-    def unit_of_measurement(self):
+    def unit_of_measurement(self) -> str:
         """Return unit of measurement."""
         return PERCENTAGE
 
     @property
-    def icon(self):
+    def icon(self) -> str:
         """Return the icon to use in the frontend, if any."""
         return icon_for_gauge_level(self.state, 10)
 
 
 class LitterRobotSleepTimeSensor(LitterRobotPropertySensor):
-    """Litter-Robot sleep time sensors."""
+    """Litter-Robot sleep time sensor."""
 
     @property
-    def state(self):
+    def state(self) -> str | None:
         """Return the state."""
-        if self.robot.sleep_mode_active:
+        if self.robot.sleep_mode_enabled:
             return super().state.isoformat()
         return None
 
     @property
-    def device_class(self):
+    def device_class(self) -> str:
         """Return the device class, if any."""
         return DEVICE_CLASS_TIMESTAMP
 
 
-ROBOT_SENSORS = [
-    (LitterRobotWasteSensor, "Waste Drawer", "waste_drawer_gauge"),
+ROBOT_SENSORS: list[tuple[type[LitterRobotPropertySensor], str, str]] = [
+    (LitterRobotWasteSensor, "Waste Drawer", "waste_drawer_level"),
     (LitterRobotSleepTimeSensor, "Sleep Mode Start Time", "sleep_mode_start_time"),
     (LitterRobotSleepTimeSensor, "Sleep Mode End Time", "sleep_mode_end_time"),
 ]
 
 
-async def async_setup_entry(hass, config_entry, async_add_entities):
+async def async_setup_entry(
+    hass: HomeAssistant,
+    entry: ConfigEntry,
+    async_add_entities: Callable[[list[Entity], bool], None],
+) -> None:
     """Set up Litter-Robot sensors using config entry."""
-    hub = hass.data[DOMAIN][config_entry.entry_id]
+    hub: LitterRobotHub = hass.data[DOMAIN][entry.entry_id]
 
     entities = []
     for robot in hub.account.robots:
         for (sensor_class, entity_type, sensor_attribute) in ROBOT_SENSORS:
-            entities.append(sensor_class(robot, entity_type, hub, sensor_attribute))
+            entities.append(
+                sensor_class(
+                    robot=robot,
+                    entity_type=entity_type,
+                    hub=hub,
+                    sensor_attribute=sensor_attribute,
+                )
+            )
 
-    if entities:
-        async_add_entities(entities, True)
+    async_add_entities(entities, True)

homeassistant/components/litterrobot/services.yaml

@@ -0,0 +1,48 @@
+# Describes the format for available Litter-Robot services
+
+reset_waste_drawer:
+  name: Reset waste drawer
+  description: Reset the waste drawer level.
+  target:
+
+set_sleep_mode:
+  name: Set sleep mode
+  description: Set the sleep mode and start time.
+  target:
+  fields:
+    enabled:
+      name: Enabled
+      description: Whether sleep mode should be enabled.
+      required: true
+      example: true
+      selector:
+        boolean:
+    start_time:
+      name: Start time
+      description: The start time at which the Litter-Robot will enter sleep mode and prevent an automatic clean cycle for 8 hours.
+      required: false
+      example: '"22:30:00"'
+      selector:
+        time:
+
+set_wait_time:
+  name: Set wait time
+  description: Set the wait time, in minutes, between when your cat uses the Litter-Robot and when the unit cycles automatically.
+  target:
+  fields:
+    minutes:
+      name: Minutes
+      description: Minutes to wait.
+      required: true
+      example: 7
+      values:
+        - 3
+        - 7
+        - 15
+      default: 7
+      selector:
+        select:
+          options:
+            - "3"
+            - "7"
+            - "15"

homeassistant/components/litterrobot/strings.json

@@ -14,7 +14,7 @@
       "unknown": "[%key:common::config_flow::error::unknown%]"
     },
     "abort": {
-      "already_configured": "[%key:common::config_flow::abort::already_configured_device%]"
+      "already_configured": "[%key:common::config_flow::abort::already_configured_account%]"
     }
   }
 }

homeassistant/components/litterrobot/switch.py

@@ -1,68 +1,79 @@
 """Support for Litter-Robot switches."""
+from __future__ import annotations
+
+from typing import Any, Callable
+
 from homeassistant.components.switch import SwitchEntity
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity import Entity
 
 from .const import DOMAIN
-from .hub import LitterRobotEntity
+from .entity import LitterRobotControlEntity
+from .hub import LitterRobotHub
 
 
-class LitterRobotNightLightModeSwitch(LitterRobotEntity, SwitchEntity):
+class LitterRobotNightLightModeSwitch(LitterRobotControlEntity, SwitchEntity):
     """Litter-Robot Night Light Mode Switch."""
 
     @property
-    def is_on(self):
+    def is_on(self) -> bool:
         """Return true if switch is on."""
-        return self.robot.night_light_active
+        return self.robot.night_light_mode_enabled
 
     @property
-    def icon(self):
+    def icon(self) -> str:
         """Return the icon."""
         return "mdi:lightbulb-on" if self.is_on else "mdi:lightbulb-off"
 
-    async def async_turn_on(self, **kwargs):
+    async def async_turn_on(self, **kwargs: Any) -> None:
         """Turn the switch on."""
         await self.perform_action_and_refresh(self.robot.set_night_light, True)
 
-    async def async_turn_off(self, **kwargs):
+    async def async_turn_off(self, **kwargs: Any) -> None:
         """Turn the switch off."""
         await self.perform_action_and_refresh(self.robot.set_night_light, False)
 
 
-class LitterRobotPanelLockoutSwitch(LitterRobotEntity, SwitchEntity):
+class LitterRobotPanelLockoutSwitch(LitterRobotControlEntity, SwitchEntity):
     """Litter-Robot Panel Lockout Switch."""
 
     @property
-    def is_on(self):
+    def is_on(self) -> bool:
         """Return true if switch is on."""
-        return self.robot.panel_lock_active
+        return self.robot.panel_lock_enabled
 
     @property
-    def icon(self):
+    def icon(self) -> str:
         """Return the icon."""
         return "mdi:lock" if self.is_on else "mdi:lock-open"
 
-    async def async_turn_on(self, **kwargs):
+    async def async_turn_on(self, **kwargs: Any) -> None:
         """Turn the switch on."""
         await self.perform_action_and_refresh(self.robot.set_panel_lockout, True)
 
-    async def async_turn_off(self, **kwargs):
+    async def async_turn_off(self, **kwargs: Any) -> None:
         """Turn the switch off."""
         await self.perform_action_and_refresh(self.robot.set_panel_lockout, False)
 
 
-ROBOT_SWITCHES = {
-    "Night Light Mode": LitterRobotNightLightModeSwitch,
-    "Panel Lockout": LitterRobotPanelLockoutSwitch,
-}
+ROBOT_SWITCHES: list[tuple[type[LitterRobotControlEntity], str]] = [
+    (LitterRobotNightLightModeSwitch, "Night Light Mode"),
+    (LitterRobotPanelLockoutSwitch, "Panel Lockout"),
+]
 
 
-async def async_setup_entry(hass, config_entry, async_add_entities):
+async def async_setup_entry(
+    hass: HomeAssistant,
+    entry: ConfigEntry,
+    async_add_entities: Callable[[list[Entity], bool], None],
+) -> None:
     """Set up Litter-Robot switches using config entry."""
-    hub = hass.data[DOMAIN][config_entry.entry_id]
+    hub: LitterRobotHub = hass.data[DOMAIN][entry.entry_id]
 
     entities = []
     for robot in hub.account.robots:
-        for switch_type, switch_class in ROBOT_SWITCHES.items():
-            entities.append(switch_class(robot, switch_type, hub))
+        for switch_class, switch_type in ROBOT_SWITCHES:
+            entities.append(switch_class(robot=robot, entity_type=switch_type, hub=hub))
 
-    if entities:
-        async_add_entities(entities, True)
+    async_add_entities(entities, True)

homeassistant/components/litterrobot/translations/en.json

@@ -1,7 +1,7 @@
 {
     "config": {
         "abort": {
-            "already_configured": "Device is already configured"
+            "already_configured": "Account is already configured"
         },
         "error": {
             "cannot_connect": "Failed to connect",

homeassistant/components/litterrobot/vacuum.py

@@ -1,11 +1,17 @@
 """Support for Litter-Robot "Vacuum"."""
-from pylitterbot import Robot
+from __future__ import annotations
+
+from typing import Any, Callable
+
+from pylitterbot.enums import LitterBoxStatus
+from pylitterbot.robot import VALID_WAIT_TIMES
+import voluptuous as vol
 
 from homeassistant.components.vacuum import (
     STATE_CLEANING,
     STATE_DOCKED,
     STATE_ERROR,
-    SUPPORT_SEND_COMMAND,
+    STATE_PAUSED,
     SUPPORT_START,
     SUPPORT_STATE,
     SUPPORT_STATUS,
@@ -13,111 +19,134 @@ from homeassistant.components.vacuum import (
     SUPPORT_TURN_ON,
     VacuumEntity,
 )
+from homeassistant.config_entries import ConfigEntry
 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 import Entity
 
 from .const import DOMAIN
-from .hub import LitterRobotEntity
+from .entity import LitterRobotControlEntity
+from .hub import LitterRobotHub
 
 SUPPORT_LITTERROBOT = (
-    SUPPORT_SEND_COMMAND
-    | SUPPORT_START
-    | SUPPORT_STATE
-    | SUPPORT_STATUS
-    | SUPPORT_TURN_OFF
-    | SUPPORT_TURN_ON
+    SUPPORT_START | SUPPORT_STATE | SUPPORT_STATUS | SUPPORT_TURN_OFF | SUPPORT_TURN_ON
 )
 TYPE_LITTER_BOX = "Litter Box"
 
+SERVICE_RESET_WASTE_DRAWER = "reset_waste_drawer"
+SERVICE_SET_SLEEP_MODE = "set_sleep_mode"
+SERVICE_SET_WAIT_TIME = "set_wait_time"
 
-async def async_setup_entry(hass, config_entry, async_add_entities):
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    entry: ConfigEntry,
+    async_add_entities: Callable[[list[Entity], bool], None],
+) -> None:
     """Set up Litter-Robot cleaner using config entry."""
-    hub = hass.data[DOMAIN][config_entry.entry_id]
+    hub: LitterRobotHub = hass.data[DOMAIN][entry.entry_id]
 
     entities = []
     for robot in hub.account.robots:
-        entities.append(LitterRobotCleaner(robot, TYPE_LITTER_BOX, hub))
+        entities.append(
+            LitterRobotCleaner(robot=robot, entity_type=TYPE_LITTER_BOX, hub=hub)
+        )
 
-    if entities:
-        async_add_entities(entities, True)
+    async_add_entities(entities, True)
+
+    platform = entity_platform.current_platform.get()
+    platform.async_register_entity_service(
+        SERVICE_RESET_WASTE_DRAWER,
+        {},
+        "async_reset_waste_drawer",
+    )
+    platform.async_register_entity_service(
+        SERVICE_SET_SLEEP_MODE,
+        {
+            vol.Required("enabled"): cv.boolean,
+            vol.Optional("start_time"): cv.time,
+        },
+        "async_set_sleep_mode",
+    )
+    platform.async_register_entity_service(
+        SERVICE_SET_WAIT_TIME,
+        {vol.Required("minutes"): vol.In(VALID_WAIT_TIMES)},
+        "async_set_wait_time",
+    )
 
 
-class LitterRobotCleaner(LitterRobotEntity, VacuumEntity):
+class LitterRobotCleaner(LitterRobotControlEntity, VacuumEntity):
     """Litter-Robot "Vacuum" Cleaner."""
 
     @property
-    def supported_features(self):
+    def supported_features(self) -> int:
         """Flag cleaner robot features that are supported."""
         return SUPPORT_LITTERROBOT
 
     @property
-    def state(self):
+    def state(self) -> str:
         """Return the state of the cleaner."""
         switcher = {
-            Robot.UnitStatus.CLEAN_CYCLE: STATE_CLEANING,
-            Robot.UnitStatus.EMPTY_CYCLE: STATE_CLEANING,
-            Robot.UnitStatus.CLEAN_CYCLE_COMPLETE: STATE_DOCKED,
-            Robot.UnitStatus.CAT_SENSOR_TIMING: STATE_DOCKED,
-            Robot.UnitStatus.DRAWER_FULL_1: STATE_DOCKED,
-            Robot.UnitStatus.DRAWER_FULL_2: STATE_DOCKED,
-            Robot.UnitStatus.READY: STATE_DOCKED,
-            Robot.UnitStatus.OFF: STATE_OFF,
+            LitterBoxStatus.CLEAN_CYCLE: STATE_CLEANING,
+            LitterBoxStatus.EMPTY_CYCLE: STATE_CLEANING,
+            LitterBoxStatus.CLEAN_CYCLE_COMPLETE: STATE_DOCKED,
+            LitterBoxStatus.CAT_SENSOR_TIMING: STATE_DOCKED,
+            LitterBoxStatus.DRAWER_FULL_1: STATE_DOCKED,
+            LitterBoxStatus.DRAWER_FULL_2: STATE_DOCKED,
+            LitterBoxStatus.READY: STATE_DOCKED,
+            LitterBoxStatus.CAT_SENSOR_INTERRUPTED: STATE_PAUSED,
+            LitterBoxStatus.OFF: STATE_OFF,
         }
 
-        return switcher.get(self.robot.unit_status, STATE_ERROR)
+        return switcher.get(self.robot.status, STATE_ERROR)
 
     @property
-    def status(self):
+    def status(self) -> str:
         """Return the status of the cleaner."""
-        return f"{self.robot.unit_status.label}{' (Sleeping)' if self.robot.is_sleeping else ''}"
+        return (
+            f"{self.robot.status.text}{' (Sleeping)' if self.robot.is_sleeping else ''}"
+        )
 
-    async def async_turn_on(self, **kwargs):
+    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)
 
-    async def async_turn_off(self, **kwargs):
+    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)
 
-    async def async_start(self):
+    async def async_start(self) -> None:
         """Start a clean cycle."""
         await self.perform_action_and_refresh(self.robot.start_cleaning)
 
-    async def async_send_command(self, command, params=None, **kwargs):
-        """Send command.
+    async def async_reset_waste_drawer(self) -> None:
+        """Reset the waste drawer level."""
+        await self.robot.reset_waste_drawer()
+        self.coordinator.async_set_updated_data(True)
 
-        Available commands:
-          - reset_waste_drawer
-            * params: none
-          - set_sleep_mode
-            * params:
-              - enabled: bool
-              - sleep_time: str (optional)
+    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(
+            self.robot.set_sleep_mode,
+            enabled,
+            self.parse_time_at_default_timezone(start_time),
+        )
 
-        """
-        if command == "reset_waste_drawer":
-            # Normally we need to request a refresh of data after a command is sent.
-            # However, the API for resetting the waste drawer returns a refreshed
-            # data set for the robot. Thus, we only need to tell hass to update the
-            # state of devices associated with this robot.
-            await self.robot.reset_waste_drawer()
-            self.hub.coordinator.async_set_updated_data(True)
-        elif command == "set_sleep_mode":
-            await self.perform_action_and_refresh(
-                self.robot.set_sleep_mode,
-                params.get("enabled"),
-                self.parse_time_at_default_timezone(params.get("sleep_time")),
-            )
-        else:
-            raise NotImplementedError()
+    async def async_set_wait_time(self, minutes: int) -> None:
+        """Set the wait time."""
+        await self.perform_action_and_refresh(self.robot.set_wait_time, minutes)
 
     @property
-    def extra_state_attributes(self):
+    def extra_state_attributes(self) -> dict[str, Any]:
         """Return device specific state attributes."""
         return {
             "clean_cycle_wait_time_minutes": self.robot.clean_cycle_wait_time_minutes,
             "is_sleeping": self.robot.is_sleeping,
-            "sleep_mode_active": self.robot.sleep_mode_active,
+            "sleep_mode_enabled": self.robot.sleep_mode_enabled,
             "power_status": self.robot.power_status,
-            "unit_status_code": self.robot.unit_status.value,
+            "status_code": self.robot.status_code,
             "last_seen": self.robot.last_seen,
         }

requirements_all.txt

@@ -1515,7 +1515,7 @@ pylibrespot-java==0.1.0
 pylitejet==0.3.0
 
 # homeassistant.components.litterrobot
-pylitterbot==2021.2.8
+pylitterbot==2021.3.1
 
 # homeassistant.components.loopenergy
 pyloopenergy==0.2.1

requirements_test_all.txt

@@ -826,7 +826,7 @@ pylibrespot-java==0.1.0
 pylitejet==0.3.0
 
 # homeassistant.components.litterrobot
-pylitterbot==2021.2.8
+pylitterbot==2021.3.1
 
 # homeassistant.components.lutron_caseta
 pylutron-caseta==0.9.0

tests/components/litterrobot/common.py

@@ -1,4 +1,6 @@
 """Common utils for Litter-Robot tests."""
+from datetime import datetime
+
 from homeassistant.components.litterrobot import DOMAIN
 from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
 
@@ -9,7 +11,7 @@ ROBOT_NAME = "Test"
 ROBOT_SERIAL = "LR3C012345"
 ROBOT_DATA = {
     "powerStatus": "AC",
-    "lastSeen": "2021-02-01T15:30:00.000000",
+    "lastSeen": datetime.now().isoformat(),
     "cleanCycleWaitTimeMinutes": "7",
     "unitStatus": "RDY",
     "litterRobotNickname": ROBOT_NAME,
@@ -22,3 +24,5 @@ ROBOT_DATA = {
     "nightLightActive": "1",
     "sleepModeActive": "112:50:19",
 }
+
+VACUUM_ENTITY_ID = "vacuum.test_litter_box"

tests/components/litterrobot/conftest.py

@@ -1,60 +1,79 @@
 """Configure pytest for Litter-Robot tests."""
-from __future__ import annotations
-
+from typing import Any, Optional
 from unittest.mock import AsyncMock, MagicMock, patch
 
-import pylitterbot
-from pylitterbot import Robot
+from pylitterbot import Account, Robot
+from pylitterbot.exceptions import InvalidCommandException
 import pytest
 
 from homeassistant.components import litterrobot
+from homeassistant.core import HomeAssistant
 
 from .common import CONFIG, ROBOT_DATA
 
 from tests.common import MockConfigEntry
 
 
-def create_mock_robot(unit_status_code: str | None = None):
+def create_mock_robot(
+    robot_data: Optional[dict] = None, side_effect: Optional[Any] = None
+) -> Robot:
     """Create a mock Litter-Robot device."""
-    if not (
-        unit_status_code
-        and Robot.UnitStatus(unit_status_code) != Robot.UnitStatus.UNKNOWN
-    ):
-        unit_status_code = ROBOT_DATA["unitStatus"]
+    if not robot_data:
+        robot_data = {}
 
-    with patch.dict(ROBOT_DATA, {"unitStatus": unit_status_code}):
-        robot = Robot(data=ROBOT_DATA)
-        robot.start_cleaning = AsyncMock()
-        robot.set_power_status = AsyncMock()
-        robot.reset_waste_drawer = AsyncMock()
-        robot.set_sleep_mode = AsyncMock()
-        robot.set_night_light = AsyncMock()
-        robot.set_panel_lockout = AsyncMock()
-        return robot
+    robot = Robot(data={**ROBOT_DATA, **robot_data})
+    robot.start_cleaning = AsyncMock(side_effect=side_effect)
+    robot.set_power_status = AsyncMock(side_effect=side_effect)
+    robot.reset_waste_drawer = AsyncMock(side_effect=side_effect)
+    robot.set_sleep_mode = AsyncMock(side_effect=side_effect)
+    robot.set_night_light = AsyncMock(side_effect=side_effect)
+    robot.set_panel_lockout = AsyncMock(side_effect=side_effect)
+    robot.set_wait_time = AsyncMock(side_effect=side_effect)
+    return robot
 
 
-def create_mock_account(unit_status_code: str | None = None):
+def create_mock_account(
+    robot_data: Optional[dict] = None,
+    side_effect: Optional[Any] = None,
+    skip_robots: bool = False,
+) -> MagicMock:
     """Create a mock Litter-Robot account."""
-    account = MagicMock(spec=pylitterbot.Account)
+    account = MagicMock(spec=Account)
     account.connect = AsyncMock()
     account.refresh_robots = AsyncMock()
-    account.robots = [create_mock_robot(unit_status_code)]
+    account.robots = [] if skip_robots else [create_mock_robot(robot_data, side_effect)]
     return account
 
 
 @pytest.fixture
-def mock_account():
+def mock_account() -> MagicMock:
     """Mock a Litter-Robot account."""
     return create_mock_account()
 
 
 @pytest.fixture
-def mock_account_with_error():
+def mock_account_with_no_robots() -> MagicMock:
+    """Mock a Litter-Robot account."""
+    return create_mock_account(skip_robots=True)
+
+
+@pytest.fixture
+def mock_account_with_error() -> MagicMock:
     """Mock a Litter-Robot account with error."""
-    return create_mock_account("BR")
+    return create_mock_account({"unitStatus": "BR"})
 
 
-async def setup_integration(hass, mock_account, platform_domain=None):
+@pytest.fixture
+def mock_account_with_side_effects() -> MagicMock:
+    """Mock a Litter-Robot account with side effects."""
+    return create_mock_account(
+        side_effect=InvalidCommandException("Invalid command: oops")
+    )
+
+
+async def setup_integration(
+    hass: HomeAssistant, mock_account: MagicMock, platform_domain: Optional[str] = None
+) -> MockConfigEntry:
     """Load a Litter-Robot platform with the provided hub."""
     entry = MockConfigEntry(
         domain=litterrobot.DOMAIN,
@@ -62,7 +81,9 @@ async def setup_integration(hass, mock_account, platform_domain=None):
     )
     entry.add_to_hass(hass)
 
-    with patch("pylitterbot.Account", return_value=mock_account), patch(
+    with patch(
+        "homeassistant.components.litterrobot.hub.Account", return_value=mock_account
+    ), patch(
         "homeassistant.components.litterrobot.PLATFORMS",
         [platform_domain] if platform_domain else [],
     ):

tests/components/litterrobot/test_config_flow.py

@@ -20,7 +20,10 @@ async def test_form(hass, mock_account):
     assert result["type"] == "form"
     assert result["errors"] == {}
 
-    with patch("pylitterbot.Account", return_value=mock_account), patch(
+    with patch(
+        "homeassistant.components.litterrobot.hub.Account",
+        return_value=mock_account,
+    ), patch(
         "homeassistant.components.litterrobot.async_setup", return_value=True
     ) as mock_setup, patch(
         "homeassistant.components.litterrobot.async_setup_entry",

tests/components/litterrobot/test_init.py

@@ -5,12 +5,18 @@ from pylitterbot.exceptions import LitterRobotException, LitterRobotLoginExcepti
 import pytest
 
 from homeassistant.components import litterrobot
+from homeassistant.components.vacuum import (
+    DOMAIN as VACUUM_DOMAIN,
+    SERVICE_START,
+    STATE_DOCKED,
+)
 from homeassistant.config_entries import (
     ENTRY_STATE_SETUP_ERROR,
     ENTRY_STATE_SETUP_RETRY,
 )
+from homeassistant.const import ATTR_ENTITY_ID
 
-from .common import CONFIG
+from .common import CONFIG, VACUUM_ENTITY_ID
 from .conftest import setup_integration
 
 from tests.common import MockConfigEntry
@@ -18,7 +24,19 @@ from tests.common import MockConfigEntry
 
 async def test_unload_entry(hass, mock_account):
     """Test being able to unload an entry."""
-    entry = await setup_integration(hass, mock_account)
+    entry = await setup_integration(hass, mock_account, VACUUM_DOMAIN)
+
+    vacuum = hass.states.get(VACUUM_ENTITY_ID)
+    assert vacuum
+    assert vacuum.state == STATE_DOCKED
+
+    await hass.services.async_call(
+        VACUUM_DOMAIN,
+        SERVICE_START,
+        {ATTR_ENTITY_ID: VACUUM_ENTITY_ID},
+        blocking=True,
+    )
+    getattr(mock_account.robots[0], "start_cleaning").assert_called_once()
 
     assert await hass.config_entries.async_unload(entry.entry_id)
     await hass.async_block_till_done()

tests/components/litterrobot/test_sensor.py

@@ -16,14 +16,13 @@ async def test_waste_drawer_sensor(hass, mock_account):
 
     sensor = hass.states.get(WASTE_DRAWER_ENTITY_ID)
     assert sensor
-    assert sensor.state == "50"
+    assert sensor.state == "50.0"
     assert sensor.attributes["unit_of_measurement"] == PERCENTAGE
 
 
 async def test_sleep_time_sensor_with_none_state(hass):
     """Tests the sleep mode start time sensor where sleep mode is inactive."""
-    robot = create_mock_robot()
-    robot.sleep_mode_active = False
+    robot = create_mock_robot({"sleepModeActive": "0"})
     sensor = LitterRobotSleepTimeSensor(
         robot, "Sleep Mode Start Time", Mock(), "sleep_mode_start_time"
     )

tests/components/litterrobot/test_switch.py

@@ -3,7 +3,7 @@ from datetime import timedelta
 
 import pytest
 
-from homeassistant.components.litterrobot.hub import REFRESH_WAIT_TIME
+from homeassistant.components.litterrobot.entity import REFRESH_WAIT_TIME_SECONDS
 from homeassistant.components.switch import (
     DOMAIN as PLATFORM_DOMAIN,
     SERVICE_TURN_OFF,
@@ -56,6 +56,6 @@ async def test_on_off_commands(hass, mock_account, entity_id, robot_command):
             blocking=True,
         )
 
-        future = utcnow() + timedelta(seconds=REFRESH_WAIT_TIME)
+        future = utcnow() + timedelta(seconds=REFRESH_WAIT_TIME_SECONDS)
         async_fire_time_changed(hass, future)
         assert getattr(mock_account.robots[0], robot_command).call_count == count

tests/components/litterrobot/test_vacuum.py

@@ -3,42 +3,60 @@ from datetime import timedelta
 
 import pytest
 
-from homeassistant.components.litterrobot.hub import REFRESH_WAIT_TIME
+from homeassistant.components.litterrobot import DOMAIN
+from homeassistant.components.litterrobot.entity import REFRESH_WAIT_TIME_SECONDS
+from homeassistant.components.litterrobot.vacuum import (
+    SERVICE_RESET_WASTE_DRAWER,
+    SERVICE_SET_SLEEP_MODE,
+    SERVICE_SET_WAIT_TIME,
+)
 from homeassistant.components.vacuum import (
-    ATTR_PARAMS,
     DOMAIN as PLATFORM_DOMAIN,
-    SERVICE_SEND_COMMAND,
     SERVICE_START,
     SERVICE_TURN_OFF,
     SERVICE_TURN_ON,
     STATE_DOCKED,
     STATE_ERROR,
 )
-from homeassistant.const import ATTR_COMMAND, ATTR_ENTITY_ID
+from homeassistant.const import ATTR_ENTITY_ID
+from homeassistant.core import HomeAssistant
 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
 
-ENTITY_ID = "vacuum.test_litter_box"
+COMPONENT_SERVICE_DOMAIN = {
+    SERVICE_RESET_WASTE_DRAWER: DOMAIN,
+    SERVICE_SET_SLEEP_MODE: DOMAIN,
+    SERVICE_SET_WAIT_TIME: DOMAIN,
+}
 
 
-async def test_vacuum(hass, mock_account):
+async def test_vacuum(hass: HomeAssistant, mock_account):
     """Tests the vacuum entity was set up."""
     await setup_integration(hass, mock_account, PLATFORM_DOMAIN)
+    assert hass.services.has_service(DOMAIN, SERVICE_RESET_WASTE_DRAWER)
 
-    vacuum = hass.states.get(ENTITY_ID)
+    vacuum = hass.states.get(VACUUM_ENTITY_ID)
     assert vacuum
     assert vacuum.state == STATE_DOCKED
     assert vacuum.attributes["is_sleeping"] is False
 
 
-async def test_vacuum_with_error(hass, mock_account_with_error):
+async def test_no_robots(hass: HomeAssistant, mock_account_with_no_robots):
+    """Tests the vacuum entity was set up."""
+    await setup_integration(hass, mock_account_with_no_robots, PLATFORM_DOMAIN)
+
+    assert not hass.services.has_service(DOMAIN, SERVICE_RESET_WASTE_DRAWER)
+
+
+async def test_vacuum_with_error(hass: HomeAssistant, mock_account_with_error):
     """Tests a vacuum entity with an error."""
     await setup_integration(hass, mock_account_with_error, PLATFORM_DOMAIN)
 
-    vacuum = hass.states.get(ENTITY_ID)
+    vacuum = hass.states.get(VACUUM_ENTITY_ID)
     assert vacuum
     assert vacuum.state == STATE_ERROR
 
@@ -50,46 +68,70 @@ async def test_vacuum_with_error(hass, mock_account_with_error):
         (SERVICE_TURN_OFF, "set_power_status", None),
         (SERVICE_TURN_ON, "set_power_status", None),
         (
-            SERVICE_SEND_COMMAND,
+            SERVICE_RESET_WASTE_DRAWER,
             "reset_waste_drawer",
-            {ATTR_COMMAND: "reset_waste_drawer"},
+            None,
         ),
         (
-            SERVICE_SEND_COMMAND,
+            SERVICE_SET_SLEEP_MODE,
             "set_sleep_mode",
-            {
-                ATTR_COMMAND: "set_sleep_mode",
-                ATTR_PARAMS: {"enabled": True, "sleep_time": "22:30"},
-            },
+            {"enabled": True, "start_time": "22:30"},
         ),
         (
-            SERVICE_SEND_COMMAND,
+            SERVICE_SET_SLEEP_MODE,
             "set_sleep_mode",
-            {
-                ATTR_COMMAND: "set_sleep_mode",
-                ATTR_PARAMS: {"enabled": True, "sleep_time": None},
-            },
+            {"enabled": True},
+        ),
+        (
+            SERVICE_SET_SLEEP_MODE,
+            "set_sleep_mode",
+            {"enabled": False},
+        ),
+        (
+            SERVICE_SET_WAIT_TIME,
+            "set_wait_time",
+            {"minutes": 3},
         ),
     ],
 )
-async def test_commands(hass, mock_account, service, command, extra):
+async def test_commands(hass: HomeAssistant, mock_account, service, command, extra):
     """Test sending commands to the vacuum."""
     await setup_integration(hass, mock_account, PLATFORM_DOMAIN)
 
-    vacuum = hass.states.get(ENTITY_ID)
+    vacuum = hass.states.get(VACUUM_ENTITY_ID)
     assert vacuum
     assert vacuum.state == STATE_DOCKED
 
-    data = {ATTR_ENTITY_ID: ENTITY_ID}
+    data = {ATTR_ENTITY_ID: VACUUM_ENTITY_ID}
     if extra:
         data.update(extra)
 
     await hass.services.async_call(
-        PLATFORM_DOMAIN,
+        COMPONENT_SERVICE_DOMAIN.get(service, PLATFORM_DOMAIN),
         service,
         data,
         blocking=True,
     )
-    future = utcnow() + timedelta(seconds=REFRESH_WAIT_TIME)
+    future = utcnow() + timedelta(seconds=REFRESH_WAIT_TIME_SECONDS)
     async_fire_time_changed(hass, future)
     getattr(mock_account.robots[0], command).assert_called_once()
+
+
+async def test_invalid_commands(
+    hass: HomeAssistant, caplog, mock_account_with_side_effects
+):
+    """Test sending invalid commands to the vacuum."""
+    await setup_integration(hass, mock_account_with_side_effects, PLATFORM_DOMAIN)
+
+    vacuum = hass.states.get(VACUUM_ENTITY_ID)
+    assert vacuum
+    assert vacuum.state == STATE_DOCKED
+
+    await hass.services.async_call(
+        DOMAIN,
+        SERVICE_SET_WAIT_TIME,
+        {ATTR_ENTITY_ID: VACUUM_ENTITY_ID, "minutes": 15},
+        blocking=True,
+    )
+    mock_account_with_side_effects.robots[0].set_wait_time.assert_called_once()
+    assert "Invalid command: oops" in caplog.text