Adjust litterrobot tests and code to match guidelines (#47060)

* Use SwitchEntity instead of ToggleEntity and adjust test patches as recommended

* Move async_create_entry out of try block in config_flow

* Patch pypi package instead of HA code

* Bump pylitterbot to 2021.2.6, fix tests, and implement other code review suggestions

* Bump pylitterbot to 2021.2.8, remove sleep mode start/end time from vacuum, adjust and add sensors for sleep mode start/end time

* Move icon helper back to Litter-Robot component and isoformat times on time sensors

homeassistant/components/litterrobot/config_flow.py

@@ -35,9 +35,6 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
             hub = LitterRobotHub(self.hass, user_input)
             try:
                 await hub.login()
-                return self.async_create_entry(
-                    title=user_input[CONF_USERNAME], data=user_input
-                )
             except LitterRobotLoginException:
                 errors["base"] = "invalid_auth"
             except LitterRobotException:
@@ -46,6 +43,11 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
                 _LOGGER.exception("Unexpected exception")
                 errors["base"] = "unknown"
 
+            if not errors:
+                return self.async_create_entry(
+                    title=user_input[CONF_USERNAME], data=user_input
+                )
+
         return self.async_show_form(
             step_id="user", data_schema=STEP_USER_DATA_SCHEMA, errors=errors
         )

homeassistant/components/litterrobot/hub.py

@@ -4,7 +4,7 @@ import logging
 from types import MethodType
 from typing import Any, Optional
 
-from pylitterbot import Account, Robot
+import pylitterbot
 from pylitterbot.exceptions import LitterRobotException, LitterRobotLoginException
 
 from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
@@ -49,7 +49,7 @@ class LitterRobotHub:
     async def login(self, load_robots: bool = False):
         """Login to Litter-Robot."""
         self.logged_in = False
-        self.account = Account()
+        self.account = pylitterbot.Account()
         try:
             await self.account.connect(
                 username=self._data[CONF_USERNAME],
@@ -69,11 +69,11 @@ class LitterRobotHub:
 class LitterRobotEntity(CoordinatorEntity):
     """Generic Litter-Robot entity representing common data and methods."""
 
-    def __init__(self, robot: Robot, entity_type: str, hub: LitterRobotHub):
+    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 if entity_type else ""
+        self.entity_type = entity_type
         self.hub = hub
 
     @property
@@ -89,22 +89,21 @@ class LitterRobotEntity(CoordinatorEntity):
     @property
     def device_info(self):
         """Return the device information for a Litter-Robot."""
-        model = "Litter-Robot 3 Connect"
-        if not self.robot.serial.startswith("LR3C"):
-            model = "Other Litter-Robot Connected Device"
         return {
             "identifiers": {(DOMAIN, self.robot.serial)},
             "name": self.robot.name,
             "manufacturer": "Litter-Robot",
-            "model": model,
+            "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, self.hub.coordinator.async_request_refresh
-        )
+        async_call_later(self.hass, REFRESH_WAIT_TIME, async_call_later_callback)
 
     @staticmethod
     def parse_time_at_default_timezone(time_str: str) -> Optional[time]:

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.5"],
+  "requirements": ["pylitterbot==2021.2.8"],
   "codeowners": ["@natekspencer"]
 }

homeassistant/components/litterrobot/sensor.py

@@ -1,32 +1,44 @@
 """Support for Litter-Robot sensors."""
-from homeassistant.const import PERCENTAGE
+from typing import Optional
+
+from pylitterbot.robot import Robot
+
+from homeassistant.const import DEVICE_CLASS_TIMESTAMP, PERCENTAGE
 from homeassistant.helpers.entity import Entity
 
 from .const import DOMAIN
-from .hub import LitterRobotEntity
-
-WASTE_DRAWER = "Waste Drawer"
+from .hub import LitterRobotEntity, LitterRobotHub
 
 
-async def async_setup_entry(hass, config_entry, async_add_entities):
-    """Set up Litter-Robot sensors using config entry."""
-    hub = hass.data[DOMAIN][config_entry.entry_id]
-
-    entities = []
-    for robot in hub.account.robots:
-        entities.append(LitterRobotSensor(robot, WASTE_DRAWER, hub))
-
-    if entities:
-        async_add_entities(entities, True)
+def icon_for_gauge_level(gauge_level: Optional[int] = None, offset: int = 0) -> str:
+    """Return a gauge icon valid identifier."""
+    if gauge_level is None or gauge_level <= 0 + offset:
+        return "mdi:gauge-empty"
+    if gauge_level > 70 + offset:
+        return "mdi:gauge-full"
+    if gauge_level > 30 + offset:
+        return "mdi:gauge"
+    return "mdi:gauge-low"
 
 
-class LitterRobotSensor(LitterRobotEntity, Entity):
-    """Litter-Robot sensors."""
+class LitterRobotPropertySensor(LitterRobotEntity, Entity):
+    """Litter-Robot property sensors."""
+
+    def __init__(
+        self, robot: Robot, entity_type: str, hub: LitterRobotHub, sensor_attribute: str
+    ):
+        """Pass coordinator to CoordinatorEntity."""
+        super().__init__(robot, entity_type, hub)
+        self.sensor_attribute = sensor_attribute
 
     @property
     def state(self):
         """Return the state."""
-        return self.robot.waste_drawer_gauge
+        return getattr(self.robot, self.sensor_attribute)
+
+
+class LitterRobotWasteSensor(LitterRobotPropertySensor, Entity):
+    """Litter-Robot sensors."""
 
     @property
     def unit_of_measurement(self):
@@ -36,19 +48,40 @@ class LitterRobotSensor(LitterRobotEntity, Entity):
     @property
     def icon(self):
         """Return the icon to use in the frontend, if any."""
-        if self.robot.waste_drawer_gauge <= 10:
-            return "mdi:gauge-empty"
-        if self.robot.waste_drawer_gauge < 50:
-            return "mdi:gauge-low"
-        if self.robot.waste_drawer_gauge <= 90:
-            return "mdi:gauge"
-        return "mdi:gauge-full"
+        return icon_for_gauge_level(self.state, 10)
+
+
+class LitterRobotSleepTimeSensor(LitterRobotPropertySensor, Entity):
+    """Litter-Robot sleep time sensors."""
 
     @property
-    def device_state_attributes(self):
-        """Return device specific state attributes."""
-        return {
-            "cycle_count": self.robot.cycle_count,
-            "cycle_capacity": self.robot.cycle_capacity,
-            "cycles_after_drawer_full": self.robot.cycles_after_drawer_full,
-        }
+    def state(self):
+        """Return the state."""
+        if self.robot.sleep_mode_active:
+            return super().state.isoformat()
+        return None
+
+    @property
+    def device_class(self):
+        """Return the device class, if any."""
+        return DEVICE_CLASS_TIMESTAMP
+
+
+ROBOT_SENSORS = [
+    (LitterRobotWasteSensor, "Waste Drawer", "waste_drawer_gauge"),
+    (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):
+    """Set up Litter-Robot sensors using config entry."""
+    hub = hass.data[DOMAIN][config_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))
+
+    if entities:
+        async_add_entities(entities, True)

homeassistant/components/litterrobot/switch.py

@@ -1,11 +1,11 @@
 """Support for Litter-Robot switches."""
-from homeassistant.helpers.entity import ToggleEntity
+from homeassistant.components.switch import SwitchEntity
 
 from .const import DOMAIN
 from .hub import LitterRobotEntity
 
 
-class LitterRobotNightLightModeSwitch(LitterRobotEntity, ToggleEntity):
+class LitterRobotNightLightModeSwitch(LitterRobotEntity, SwitchEntity):
     """Litter-Robot Night Light Mode Switch."""
 
     @property
@@ -27,7 +27,7 @@ class LitterRobotNightLightModeSwitch(LitterRobotEntity, ToggleEntity):
         await self.perform_action_and_refresh(self.robot.set_night_light, False)
 
 
-class LitterRobotPanelLockoutSwitch(LitterRobotEntity, ToggleEntity):
+class LitterRobotPanelLockoutSwitch(LitterRobotEntity, SwitchEntity):
     """Litter-Robot Panel Lockout Switch."""
 
     @property

homeassistant/components/litterrobot/vacuum.py

@@ -14,7 +14,6 @@ from homeassistant.components.vacuum import (
     VacuumEntity,
 )
 from homeassistant.const import STATE_OFF
-import homeassistant.util.dt as dt_util
 
 from .const import DOMAIN
 from .hub import LitterRobotEntity
@@ -54,27 +53,22 @@ class LitterRobotCleaner(LitterRobotEntity, VacuumEntity):
     def state(self):
         """Return the state of the cleaner."""
         switcher = {
-            Robot.UnitStatus.CCP: STATE_CLEANING,
-            Robot.UnitStatus.EC: STATE_CLEANING,
-            Robot.UnitStatus.CCC: STATE_DOCKED,
-            Robot.UnitStatus.CST: STATE_DOCKED,
-            Robot.UnitStatus.DF1: STATE_DOCKED,
-            Robot.UnitStatus.DF2: STATE_DOCKED,
-            Robot.UnitStatus.RDY: STATE_DOCKED,
+            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,
         }
 
         return switcher.get(self.robot.unit_status, STATE_ERROR)
 
-    @property
-    def error(self):
-        """Return the error associated with the current state, if any."""
-        return self.robot.unit_status.value
-
     @property
     def status(self):
         """Return the status of the cleaner."""
-        return f"{self.robot.unit_status.value}{' (Sleeping)' if self.robot.is_sleeping else ''}"
+        return f"{self.robot.unit_status.label}{' (Sleeping)' if self.robot.is_sleeping else ''}"
 
     async def async_turn_on(self, **kwargs):
         """Turn the cleaner on, starting a clean cycle."""
@@ -119,22 +113,11 @@ class LitterRobotCleaner(LitterRobotEntity, VacuumEntity):
     @property
     def device_state_attributes(self):
         """Return device specific state attributes."""
-        [sleep_mode_start_time, sleep_mode_end_time] = [None, None]
-
-        if self.robot.sleep_mode_active:
-            sleep_mode_start_time = dt_util.as_local(
-                self.robot.sleep_mode_start_time
-            ).strftime("%H:%M:00")
-            sleep_mode_end_time = dt_util.as_local(
-                self.robot.sleep_mode_end_time
-            ).strftime("%H:%M:00")
-
         return {
             "clean_cycle_wait_time_minutes": self.robot.clean_cycle_wait_time_minutes,
             "is_sleeping": self.robot.is_sleeping,
-            "sleep_mode_start_time": sleep_mode_start_time,
-            "sleep_mode_end_time": sleep_mode_end_time,
+            "sleep_mode_active": self.robot.sleep_mode_active,
             "power_status": self.robot.power_status,
-            "unit_status_code": self.robot.unit_status.name,
+            "unit_status_code": self.robot.unit_status.value,
             "last_seen": self.robot.last_seen,
         }

requirements_all.txt

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

requirements_test_all.txt

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

tests/components/litterrobot/conftest.py

@@ -1,45 +1,59 @@
 """Configure pytest for Litter-Robot tests."""
+from typing import Optional
 from unittest.mock import AsyncMock, MagicMock, patch
 
+import pylitterbot
 from pylitterbot import Robot
 import pytest
 
 from homeassistant.components import litterrobot
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
 from .common import CONFIG, ROBOT_DATA
 
 from tests.common import MockConfigEntry
 
 
-def create_mock_robot(hass):
+def create_mock_robot(unit_status_code: Optional[str] = None):
     """Create a mock Litter-Robot device."""
-    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
+    if not (
+        unit_status_code
+        and Robot.UnitStatus(unit_status_code) != Robot.UnitStatus.UNKNOWN
+    ):
+        unit_status_code = ROBOT_DATA["unitStatus"]
+
+    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
 
 
-@pytest.fixture()
-def mock_hub(hass):
-    """Mock a Litter-Robot hub."""
-    hub = MagicMock(
-        hass=hass,
-        account=MagicMock(),
-        logged_in=True,
-        coordinator=MagicMock(spec=DataUpdateCoordinator),
-        spec=litterrobot.LitterRobotHub,
-    )
-    hub.coordinator.last_update_success = True
-    hub.account.robots = [create_mock_robot(hass)]
-    return hub
+def create_mock_account(unit_status_code: Optional[str] = None):
+    """Create a mock Litter-Robot account."""
+    account = MagicMock(spec=pylitterbot.Account)
+    account.connect = AsyncMock()
+    account.refresh_robots = AsyncMock()
+    account.robots = [create_mock_robot(unit_status_code)]
+    return account
 
 
-async def setup_hub(hass, mock_hub, platform_domain):
+@pytest.fixture
+def mock_account():
+    """Mock a Litter-Robot account."""
+    return create_mock_account()
+
+
+@pytest.fixture
+def mock_account_with_error():
+    """Mock a Litter-Robot account with error."""
+    return create_mock_account("BR")
+
+
+async def setup_integration(hass, mock_account, platform_domain=None):
     """Load a Litter-Robot platform with the provided hub."""
     entry = MockConfigEntry(
         domain=litterrobot.DOMAIN,
@@ -47,9 +61,11 @@ async def setup_hub(hass, mock_hub, platform_domain):
     )
     entry.add_to_hass(hass)
 
-    with patch(
-        "homeassistant.components.litterrobot.LitterRobotHub",
-        return_value=mock_hub,
+    with patch("pylitterbot.Account", return_value=mock_account), patch(
+        "homeassistant.components.litterrobot.PLATFORMS",
+        [platform_domain] if platform_domain else [],
     ):
-        await hass.config_entries.async_forward_entry_setup(entry, platform_domain)
+        await hass.config_entries.async_setup(entry.entry_id)
         await hass.async_block_till_done()
+
+    return entry

tests/components/litterrobot/test_config_flow.py

@@ -4,11 +4,14 @@ from unittest.mock import patch
 from pylitterbot.exceptions import LitterRobotException, LitterRobotLoginException
 
 from homeassistant import config_entries, setup
+from homeassistant.components import litterrobot
 
 from .common import CONF_USERNAME, CONFIG, DOMAIN
 
+from tests.common import MockConfigEntry
 
-async def test_form(hass):
+
+async def test_form(hass, mock_account):
     """Test we get the form."""
     await setup.async_setup_component(hass, "persistent_notification", {})
     result = await hass.config_entries.flow.async_init(
@@ -17,10 +20,7 @@ async def test_form(hass):
     assert result["type"] == "form"
     assert result["errors"] == {}
 
-    with patch(
-        "homeassistant.components.litterrobot.config_flow.LitterRobotHub.login",
-        return_value=True,
-    ), patch(
+    with patch("pylitterbot.Account", return_value=mock_account), patch(
         "homeassistant.components.litterrobot.async_setup", return_value=True
     ) as mock_setup, patch(
         "homeassistant.components.litterrobot.async_setup_entry",
@@ -38,6 +38,23 @@ async def test_form(hass):
     assert len(mock_setup_entry.mock_calls) == 1
 
 
+async def test_already_configured(hass):
+    """Test we handle already configured."""
+    MockConfigEntry(
+        domain=litterrobot.DOMAIN,
+        data=CONFIG[litterrobot.DOMAIN],
+    ).add_to_hass(hass)
+
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        context={"source": config_entries.SOURCE_USER},
+        data=CONFIG[litterrobot.DOMAIN],
+    )
+
+    assert result["type"] == "abort"
+    assert result["reason"] == "already_configured"
+
+
 async def test_form_invalid_auth(hass):
     """Test we handle invalid auth."""
     result = await hass.config_entries.flow.async_init(
@@ -45,7 +62,7 @@ async def test_form_invalid_auth(hass):
     )
 
     with patch(
-        "homeassistant.components.litterrobot.config_flow.LitterRobotHub.login",
+        "pylitterbot.Account.connect",
         side_effect=LitterRobotLoginException,
     ):
         result2 = await hass.config_entries.flow.async_configure(
@@ -63,7 +80,7 @@ async def test_form_cannot_connect(hass):
     )
 
     with patch(
-        "homeassistant.components.litterrobot.config_flow.LitterRobotHub.login",
+        "pylitterbot.Account.connect",
         side_effect=LitterRobotException,
     ):
         result2 = await hass.config_entries.flow.async_configure(
@@ -81,7 +98,7 @@ async def test_form_unknown_error(hass):
     )
 
     with patch(
-        "homeassistant.components.litterrobot.config_flow.LitterRobotHub.login",
+        "pylitterbot.Account.connect",
         side_effect=Exception,
     ):
         result2 = await hass.config_entries.flow.async_configure(

tests/components/litterrobot/test_init.py

@@ -1,20 +1,48 @@
 """Test Litter-Robot setup process."""
+from unittest.mock import patch
+
+from pylitterbot.exceptions import LitterRobotException, LitterRobotLoginException
+import pytest
+
 from homeassistant.components import litterrobot
-from homeassistant.setup import async_setup_component
+from homeassistant.config_entries import (
+    ENTRY_STATE_SETUP_ERROR,
+    ENTRY_STATE_SETUP_RETRY,
+)
 
 from .common import CONFIG
+from .conftest import setup_integration
 
 from tests.common import MockConfigEntry
 
 
-async def test_unload_entry(hass):
+async def test_unload_entry(hass, mock_account):
     """Test being able to unload an entry."""
+    entry = await setup_integration(hass, mock_account)
+
+    assert await hass.config_entries.async_unload(entry.entry_id)
+    await hass.async_block_till_done()
+    assert hass.data[litterrobot.DOMAIN] == {}
+
+
+@pytest.mark.parametrize(
+    "side_effect,expected_state",
+    (
+        (LitterRobotLoginException, ENTRY_STATE_SETUP_ERROR),
+        (LitterRobotException, ENTRY_STATE_SETUP_RETRY),
+    ),
+)
+async def test_entry_not_setup(hass, side_effect, expected_state):
+    """Test being able to handle config entry not setup."""
     entry = MockConfigEntry(
         domain=litterrobot.DOMAIN,
         data=CONFIG[litterrobot.DOMAIN],
     )
     entry.add_to_hass(hass)
 
-    assert await async_setup_component(hass, litterrobot.DOMAIN, {}) is True
-    assert await litterrobot.async_unload_entry(hass, entry)
-    assert hass.data[litterrobot.DOMAIN] == {}
+    with patch(
+        "pylitterbot.Account.connect",
+        side_effect=side_effect,
+    ):
+        await hass.config_entries.async_setup(entry.entry_id)
+        assert entry.state == expected_state

tests/components/litterrobot/test_sensor.py

@@ -1,20 +1,57 @@
 """Test the Litter-Robot sensor entity."""
+from unittest.mock import Mock
+
+from homeassistant.components.litterrobot.sensor import LitterRobotSleepTimeSensor
 from homeassistant.components.sensor import DOMAIN as PLATFORM_DOMAIN
-from homeassistant.const import PERCENTAGE
+from homeassistant.const import DEVICE_CLASS_TIMESTAMP, PERCENTAGE
 
-from .conftest import setup_hub
+from .conftest import create_mock_robot, setup_integration
 
-ENTITY_ID = "sensor.test_waste_drawer"
+WASTE_DRAWER_ENTITY_ID = "sensor.test_waste_drawer"
 
 
-async def test_sensor(hass, mock_hub):
-    """Tests the sensor entity was set up."""
-    await setup_hub(hass, mock_hub, PLATFORM_DOMAIN)
+async def test_waste_drawer_sensor(hass, mock_account):
+    """Tests the waste drawer sensor entity was set up."""
+    await setup_integration(hass, mock_account, PLATFORM_DOMAIN)
 
-    sensor = hass.states.get(ENTITY_ID)
+    sensor = hass.states.get(WASTE_DRAWER_ENTITY_ID)
     assert sensor
     assert sensor.state == "50"
-    assert sensor.attributes["cycle_count"] == 15
-    assert sensor.attributes["cycle_capacity"] == 30
-    assert sensor.attributes["cycles_after_drawer_full"] == 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
+    sensor = LitterRobotSleepTimeSensor(
+        robot, "Sleep Mode Start Time", Mock(), "sleep_mode_start_time"
+    )
+
+    assert sensor
+    assert sensor.state is None
+    assert sensor.device_class == DEVICE_CLASS_TIMESTAMP
+
+
+async def test_gauge_icon():
+    """Test icon generator for gauge sensor."""
+    from homeassistant.components.litterrobot.sensor import icon_for_gauge_level
+
+    GAUGE_EMPTY = "mdi:gauge-empty"
+    GAUGE_LOW = "mdi:gauge-low"
+    GAUGE = "mdi:gauge"
+    GAUGE_FULL = "mdi:gauge-full"
+
+    assert icon_for_gauge_level(None) == GAUGE_EMPTY
+    assert icon_for_gauge_level(0) == GAUGE_EMPTY
+    assert icon_for_gauge_level(5) == GAUGE_LOW
+    assert icon_for_gauge_level(40) == GAUGE
+    assert icon_for_gauge_level(80) == GAUGE_FULL
+    assert icon_for_gauge_level(100) == GAUGE_FULL
+
+    assert icon_for_gauge_level(None, 10) == GAUGE_EMPTY
+    assert icon_for_gauge_level(0, 10) == GAUGE_EMPTY
+    assert icon_for_gauge_level(5, 10) == GAUGE_EMPTY
+    assert icon_for_gauge_level(40, 10) == GAUGE_LOW
+    assert icon_for_gauge_level(80, 10) == GAUGE
+    assert icon_for_gauge_level(100, 10) == GAUGE_FULL

tests/components/litterrobot/test_switch.py

@@ -12,7 +12,7 @@ from homeassistant.components.switch import (
 from homeassistant.const import ATTR_ENTITY_ID, STATE_ON
 from homeassistant.util.dt import utcnow
 
-from .conftest import setup_hub
+from .conftest import setup_integration
 
 from tests.common import async_fire_time_changed
 
@@ -20,9 +20,9 @@ NIGHT_LIGHT_MODE_ENTITY_ID = "switch.test_night_light_mode"
 PANEL_LOCKOUT_ENTITY_ID = "switch.test_panel_lockout"
 
 
-async def test_switch(hass, mock_hub):
+async def test_switch(hass, mock_account):
     """Tests the switch entity was set up."""
-    await setup_hub(hass, mock_hub, PLATFORM_DOMAIN)
+    await setup_integration(hass, mock_account, PLATFORM_DOMAIN)
 
     switch = hass.states.get(NIGHT_LIGHT_MODE_ENTITY_ID)
     assert switch
@@ -36,9 +36,9 @@ async def test_switch(hass, mock_hub):
         (PANEL_LOCKOUT_ENTITY_ID, "set_panel_lockout"),
     ],
 )
-async def test_on_off_commands(hass, mock_hub, entity_id, robot_command):
+async def test_on_off_commands(hass, mock_account, entity_id, robot_command):
     """Test sending commands to the switch."""
-    await setup_hub(hass, mock_hub, PLATFORM_DOMAIN)
+    await setup_integration(hass, mock_account, PLATFORM_DOMAIN)
 
     switch = hass.states.get(entity_id)
     assert switch
@@ -48,12 +48,14 @@ async def test_on_off_commands(hass, mock_hub, entity_id, robot_command):
     count = 0
     for service in [SERVICE_TURN_ON, SERVICE_TURN_OFF]:
         count += 1
+
         await hass.services.async_call(
             PLATFORM_DOMAIN,
             service,
             data,
             blocking=True,
         )
+
         future = utcnow() + timedelta(seconds=REFRESH_WAIT_TIME)
         async_fire_time_changed(hass, future)
-        assert getattr(mock_hub.account.robots[0], robot_command).call_count == count
+        assert getattr(mock_account.robots[0], robot_command).call_count == count

tests/components/litterrobot/test_vacuum.py

@@ -12,20 +12,21 @@ from homeassistant.components.vacuum import (
     SERVICE_TURN_OFF,
     SERVICE_TURN_ON,
     STATE_DOCKED,
+    STATE_ERROR,
 )
 from homeassistant.const import ATTR_COMMAND, ATTR_ENTITY_ID
 from homeassistant.util.dt import utcnow
 
-from .conftest import setup_hub
+from .conftest import setup_integration
 
 from tests.common import async_fire_time_changed
 
 ENTITY_ID = "vacuum.test_litter_box"
 
 
-async def test_vacuum(hass, mock_hub):
+async def test_vacuum(hass, mock_account):
     """Tests the vacuum entity was set up."""
-    await setup_hub(hass, mock_hub, PLATFORM_DOMAIN)
+    await setup_integration(hass, mock_account, PLATFORM_DOMAIN)
 
     vacuum = hass.states.get(ENTITY_ID)
     assert vacuum
@@ -33,6 +34,15 @@ async def test_vacuum(hass, mock_hub):
     assert vacuum.attributes["is_sleeping"] is False
 
 
+async def test_vacuum_with_error(hass, 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)
+    assert vacuum
+    assert vacuum.state == STATE_ERROR
+
+
 @pytest.mark.parametrize(
     "service,command,extra",
     [
@@ -52,14 +62,22 @@ async def test_vacuum(hass, mock_hub):
                 ATTR_PARAMS: {"enabled": True, "sleep_time": "22:30"},
             },
         ),
+        (
+            SERVICE_SEND_COMMAND,
+            "set_sleep_mode",
+            {
+                ATTR_COMMAND: "set_sleep_mode",
+                ATTR_PARAMS: {"enabled": True, "sleep_time": None},
+            },
+        ),
     ],
 )
-async def test_commands(hass, mock_hub, service, command, extra):
+async def test_commands(hass, mock_account, service, command, extra):
     """Test sending commands to the vacuum."""
-    await setup_hub(hass, mock_hub, PLATFORM_DOMAIN)
+    await setup_integration(hass, mock_account, PLATFORM_DOMAIN)
 
     vacuum = hass.states.get(ENTITY_ID)
-    assert vacuum is not None
+    assert vacuum
     assert vacuum.state == STATE_DOCKED
 
     data = {ATTR_ENTITY_ID: ENTITY_ID}
@@ -74,4 +92,4 @@ async def test_commands(hass, mock_hub, service, command, extra):
     )
     future = utcnow() + timedelta(seconds=REFRESH_WAIT_TIME)
     async_fire_time_changed(hass, future)
-    getattr(mock_hub.account.robots[0], command).assert_called_once()
+    getattr(mock_account.robots[0], command).assert_called_once()