Add switches and sensors to Litter-Robot

homeassistant/components/litterrobot/__init__.py

@@ -10,7 +10,7 @@
 from .const import DOMAIN
 from .hub import LitterRobotHub
 
-PLATFORMS = ["vacuum"]
+PLATFORMS = ["sensor", "switch", "vacuum"]
 
 
 async def async_setup(hass: HomeAssistant, config: dict):

homeassistant/components/litterrobot/sensor.py

@@ -0,0 +1,54 @@
+"""Support for Litter-Robot sensors."""
+from homeassistant.const import PERCENTAGE
+from homeassistant.helpers.entity import Entity
+
+from .const import DOMAIN
+from .hub import LitterRobotEntity
+
+WASTE_DRAWER = "Waste Drawer"
+
+
+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)
+
+
+class LitterRobotSensor(LitterRobotEntity, Entity):
+    """Litter-Robot sensors."""
+
+    @property
+    def state(self):
+        """Return the state."""
+        return self.robot.waste_drawer_gauge
+
+    @property
+    def unit_of_measurement(self):
+        """Return unit of measurement."""
+        return PERCENTAGE
+
+    @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"
+
+    @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,
+        }

homeassistant/components/litterrobot/switch.py

@@ -0,0 +1,68 @@
+"""Support for Litter-Robot switches."""
+from homeassistant.helpers.entity import ToggleEntity
+
+from .const import DOMAIN
+from .hub import LitterRobotEntity
+
+
+class LitterRobotNightLightModeSwitch(LitterRobotEntity, ToggleEntity):
+    """Litter-Robot Night Light Mode Switch."""
+
+    @property
+    def is_on(self):
+        """Return true if switch is on."""
+        return self.robot.night_light_active
+
+    @property
+    def icon(self):
+        """Return the icon."""
+        return "mdi:lightbulb-on" if self.is_on else "mdi:lightbulb-off"
+
+    async def async_turn_on(self, **kwargs):
+        """Turn the switch on."""
+        await self.perform_action_and_refresh(self.robot.set_night_light, True)
+
+    async def async_turn_off(self, **kwargs):
+        """Turn the switch off."""
+        await self.perform_action_and_refresh(self.robot.set_night_light, False)
+
+
+class LitterRobotPanelLockoutSwitch(LitterRobotEntity, ToggleEntity):
+    """Litter-Robot Panel Lockout Switch."""
+
+    @property
+    def is_on(self):
+        """Return true if switch is on."""
+        return self.robot.panel_lock_active
+
+    @property
+    def icon(self):
+        """Return the icon."""
+        return "mdi:lock" if self.is_on else "mdi:lock-open"
+
+    async def async_turn_on(self, **kwargs):
+        """Turn the switch on."""
+        await self.perform_action_and_refresh(self.robot.set_panel_lockout, True)
+
+    async def async_turn_off(self, **kwargs):
+        """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,
+}
+
+
+async def async_setup_entry(hass, config_entry, async_add_entities):
+    """Set up Litter-Robot switches using config entry."""
+    hub = hass.data[DOMAIN][config_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))
+
+    if entities:
+        async_add_entities(entities, True)

homeassistant/components/litterrobot/vacuum.py

@@ -14,6 +14,7 @@
     VacuumEntity,
 )
 from homeassistant.const import STATE_OFF
+import homeassistant.util.dt as dt_util
 
 from .const import DOMAIN
 from .hub import LitterRobotEntity
@@ -118,9 +119,21 @@ async def async_send_command(self, command, params=None, **kwargs):
     @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,
             "power_status": self.robot.power_status,
             "unit_status_code": self.robot.unit_status.name,
             "last_seen": self.robot.last_seen,

tests/components/litterrobot/conftest.py

@@ -1,13 +1,15 @@
 """Configure pytest for Litter-Robot tests."""
-from unittest.mock import AsyncMock, MagicMock
+from unittest.mock import AsyncMock, MagicMock, patch
 
 from pylitterbot import Robot
 import pytest
 
 from homeassistant.components import litterrobot
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from .common import ROBOT_DATA
+from .common import CONFIG, ROBOT_DATA
+
+from tests.common import MockConfigEntry
 
 
 def create_mock_robot(hass):
@@ -17,6 +19,8 @@ def create_mock_robot(hass):
     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
 
 
@@ -33,3 +37,19 @@ def mock_hub(hass):
     hub.coordinator.last_update_success = True
     hub.account.robots = [create_mock_robot(hass)]
     return hub
+
+
+async def setup_hub(hass, mock_hub, platform_domain):
+    """Load a Litter-Robot platform with the provided hub."""
+    entry = MockConfigEntry(
+        domain=litterrobot.DOMAIN,
+        data=CONFIG[litterrobot.DOMAIN],
+    )
+    entry.add_to_hass(hass)
+
+    with patch(
+        "homeassistant.components.litterrobot.LitterRobotHub",
+        return_value=mock_hub,
+    ):
+        await hass.config_entries.async_forward_entry_setup(entry, platform_domain)
+        await hass.async_block_till_done()

tests/components/litterrobot/test_sensor.py

@@ -0,0 +1,20 @@
+"""Test the Litter-Robot sensor entity."""
+from homeassistant.components.sensor import DOMAIN as PLATFORM_DOMAIN
+from homeassistant.const import PERCENTAGE
+
+from .conftest import setup_hub
+
+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)
+
+    sensor = hass.states.get(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

tests/components/litterrobot/test_switch.py

@@ -0,0 +1,59 @@
+"""Test the Litter-Robot switch entity."""
+from datetime import timedelta
+
+import pytest
+
+from homeassistant.components.litterrobot.hub import REFRESH_WAIT_TIME
+from homeassistant.components.switch import (
+    DOMAIN as PLATFORM_DOMAIN,
+    SERVICE_TURN_OFF,
+    SERVICE_TURN_ON,
+)
+from homeassistant.const import ATTR_ENTITY_ID, STATE_ON
+from homeassistant.util.dt import utcnow
+
+from .conftest import setup_hub
+
+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"
+
+
+async def test_switch(hass, mock_hub):
+    """Tests the switch entity was set up."""
+    await setup_hub(hass, mock_hub, PLATFORM_DOMAIN)
+
+    switch = hass.states.get(NIGHT_LIGHT_MODE_ENTITY_ID)
+    assert switch
+    assert switch.state == STATE_ON
+
+
+@pytest.mark.parametrize(
+    "entity_id,robot_command",
+    [
+        (NIGHT_LIGHT_MODE_ENTITY_ID, "set_night_light"),
+        (PANEL_LOCKOUT_ENTITY_ID, "set_panel_lockout"),
+    ],
+)
+async def test_on_off_commands(hass, mock_hub, entity_id, robot_command):
+    """Test sending commands to the switch."""
+    await setup_hub(hass, mock_hub, PLATFORM_DOMAIN)
+
+    switch = hass.states.get(entity_id)
+    assert switch
+
+    data = {ATTR_ENTITY_ID: entity_id}
+
+    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

tests/components/litterrobot/test_vacuum.py

@@ -1,10 +1,8 @@
 """Test the Litter-Robot vacuum entity."""
 from datetime import timedelta
-from unittest.mock import patch
 
 import pytest
 
-from homeassistant.components import litterrobot
 from homeassistant.components.litterrobot.hub import REFRESH_WAIT_TIME
 from homeassistant.components.vacuum import (
     ATTR_PARAMS,
@@ -18,32 +16,19 @@
 from homeassistant.const import ATTR_COMMAND, ATTR_ENTITY_ID
 from homeassistant.util.dt import utcnow
 
-from .common import CONFIG
+from .conftest import setup_hub
 
-from tests.common import MockConfigEntry, async_fire_time_changed
+from tests.common import async_fire_time_changed
 
 ENTITY_ID = "vacuum.test_litter_box"
 
 
-async def setup_hub(hass, mock_hub):
-    """Load the Litter-Robot vacuum platform with the provided hub."""
-    hass.config.components.add(litterrobot.DOMAIN)
-    entry = MockConfigEntry(
-        domain=litterrobot.DOMAIN,
-        data=CONFIG[litterrobot.DOMAIN],
-    )
-
-    with patch.dict(hass.data, {litterrobot.DOMAIN: {entry.entry_id: mock_hub}}):
-        await hass.config_entries.async_forward_entry_setup(entry, PLATFORM_DOMAIN)
-        await hass.async_block_till_done()
-
-
 async def test_vacuum(hass, mock_hub):
     """Tests the vacuum entity was set up."""
-    await setup_hub(hass, mock_hub)
+    await setup_hub(hass, mock_hub, PLATFORM_DOMAIN)
 
     vacuum = hass.states.get(ENTITY_ID)
-    assert vacuum is not None
+    assert vacuum
     assert vacuum.state == STATE_DOCKED
     assert vacuum.attributes["is_sleeping"] is False
 
@@ -71,7 +56,7 @@ async def test_vacuum(hass, mock_hub):
 )
 async def test_commands(hass, mock_hub, service, command, extra):
     """Test sending commands to the vacuum."""
-    await setup_hub(hass, mock_hub)
+    await setup_hub(hass, mock_hub, PLATFORM_DOMAIN)
 
     vacuum = hass.states.get(ENTITY_ID)
     assert vacuum is not None