Add firmware update entity for Litter-Robot 4 (#83590)

* Add firmware update entity for Litter-Robot 4

* Report installed version of firmware on robot when updated

homeassistant/components/litterrobot/__init__.py

@@ -1,7 +1,7 @@
 """The Litter-Robot integration."""
 from __future__ import annotations
 
-from pylitterbot import FeederRobot, LitterRobot, LitterRobot3, Robot
+from pylitterbot import FeederRobot, LitterRobot, LitterRobot3, LitterRobot4, Robot
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
@@ -19,6 +19,7 @@ PLATFORMS_BY_TYPE = {
     ),
     LitterRobot: (Platform.VACUUM,),
     LitterRobot3: (Platform.BUTTON,),
+    LitterRobot4: (Platform.UPDATE,),
     FeederRobot: (Platform.BUTTON,),
 }
 

homeassistant/components/litterrobot/update.py

@@ -0,0 +1,111 @@
+"""Support for Litter-Robot updates."""
+from __future__ import annotations
+
+from collections.abc import Callable
+from datetime import datetime, timedelta
+from typing import Any
+
+from pylitterbot import LitterRobot4
+
+from homeassistant.components.update import (
+    UpdateDeviceClass,
+    UpdateEntity,
+    UpdateEntityDescription,
+    UpdateEntityFeature,
+)
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import HomeAssistantError
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.helpers.event import async_call_later
+from homeassistant.helpers.start import async_at_start
+
+from .const import DOMAIN
+from .entity import LitterRobotEntity, LitterRobotHub
+
+FIRMWARE_UPDATE_ENTITY = UpdateEntityDescription(
+    key="firmware",
+    name="Firmware",
+    device_class=UpdateDeviceClass.FIRMWARE,
+)
+
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    entry: ConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+) -> None:
+    """Set up Litter-Robot update platform."""
+    hub: LitterRobotHub = hass.data[DOMAIN][entry.entry_id]
+    robots = hub.account.robots
+    entities = [
+        RobotUpdateEntity(robot=robot, hub=hub, description=FIRMWARE_UPDATE_ENTITY)
+        for robot in robots
+        if isinstance(robot, LitterRobot4)
+    ]
+    async_add_entities(entities)
+
+
+class RobotUpdateEntity(LitterRobotEntity[LitterRobot4], UpdateEntity):
+    """A class that describes robot update entities."""
+
+    _attr_supported_features = (
+        UpdateEntityFeature.INSTALL | UpdateEntityFeature.PROGRESS
+    )
+
+    def __init__(
+        self,
+        robot: LitterRobot4,
+        hub: LitterRobotHub,
+        description: UpdateEntityDescription,
+    ) -> None:
+        """Initialize a Litter-Robot update entity."""
+        super().__init__(robot, hub, description)
+        self._poll_unsub: Callable[[], None] | None = None
+
+    @property
+    def installed_version(self) -> str:
+        """Version installed and in use."""
+        return self.robot.firmware
+
+    @property
+    def in_progress(self) -> bool:
+        """Update installation progress."""
+        return self.robot.firmware_update_triggered
+
+    async def _async_update(self, _: HomeAssistant | datetime | None = None) -> None:
+        """Update the entity."""
+        self._poll_unsub = None
+
+        if await self.robot.has_firmware_update():
+            latest_version = await self.robot.get_latest_firmware()
+        else:
+            latest_version = self.installed_version
+
+        if self._attr_latest_version != self.installed_version:
+            self._attr_latest_version = latest_version
+            self.async_write_ha_state()
+
+        self._poll_unsub = async_call_later(
+            self.hass, timedelta(days=1), self._async_update
+        )
+
+    async def async_added_to_hass(self) -> None:
+        """Set up a listener for the entity."""
+        await super().async_added_to_hass()
+        self.async_on_remove(async_at_start(self.hass, self._async_update))
+
+    async def async_install(
+        self, version: str | None, backup: bool, **kwargs: Any
+    ) -> None:
+        """Install an update."""
+        if await self.robot.has_firmware_update():
+            if not await self.robot.update_firmware():
+                message = f"Unable to start firmware update on {self.robot.name}"
+                raise HomeAssistantError(message)
+
+    async def async_will_remove_from_hass(self) -> None:
+        """Call when entity will be removed."""
+        if self._poll_unsub:
+            self._poll_unsub()
+            self._poll_unsub = None

tests/components/litterrobot/test_update.py

@@ -0,0 +1,81 @@
+"""Test the Litter-Robot update entity."""
+from unittest.mock import AsyncMock, MagicMock
+
+from pylitterbot import LitterRobot4
+import pytest
+
+from homeassistant.components.update import (
+    ATTR_INSTALLED_VERSION,
+    ATTR_LATEST_VERSION,
+    DOMAIN as PLATFORM_DOMAIN,
+    SERVICE_INSTALL,
+    UpdateDeviceClass,
+)
+from homeassistant.const import ATTR_DEVICE_CLASS, ATTR_ENTITY_ID, STATE_OFF, STATE_ON
+from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import HomeAssistantError
+
+from .conftest import setup_integration
+
+ENTITY_ID = "update.test_firmware"
+OLD_FIRMWARE = "ESP: 1.1.50 / PIC: 10512.2560.2.53 / TOF: 4.0.65.4"
+NEW_FIRMWARE = "ESP: 1.1.51 / PIC: 10512.2560.2.53 / TOF: 4.0.65.4"
+
+
+async def test_robot_with_no_update(
+    hass: HomeAssistant, mock_account_with_litterrobot_4: MagicMock
+):
+    """Tests the update entity was set up."""
+    robot: LitterRobot4 = mock_account_with_litterrobot_4.robots[0]
+    robot.has_firmware_update = AsyncMock(return_value=False)
+
+    entry = await setup_integration(
+        hass, mock_account_with_litterrobot_4, PLATFORM_DOMAIN
+    )
+
+    state = hass.states.get(ENTITY_ID)
+    assert state
+    assert state.state == STATE_OFF
+    assert state.attributes[ATTR_DEVICE_CLASS] == UpdateDeviceClass.FIRMWARE
+    assert state.attributes[ATTR_INSTALLED_VERSION] == OLD_FIRMWARE
+    assert state.attributes[ATTR_LATEST_VERSION] == OLD_FIRMWARE
+
+    assert await hass.config_entries.async_unload(entry.entry_id)
+    await hass.async_block_till_done()
+
+
+async def test_robot_with_update(
+    hass: HomeAssistant, mock_account_with_litterrobot_4: MagicMock
+):
+    """Tests the update entity was set up."""
+    robot: LitterRobot4 = mock_account_with_litterrobot_4.robots[0]
+    robot.has_firmware_update = AsyncMock(return_value=True)
+    robot.get_latest_firmware = AsyncMock(return_value=NEW_FIRMWARE)
+
+    await setup_integration(hass, mock_account_with_litterrobot_4, PLATFORM_DOMAIN)
+
+    state = hass.states.get(ENTITY_ID)
+    assert state
+    assert state.state == STATE_ON
+    assert state.attributes[ATTR_DEVICE_CLASS] == UpdateDeviceClass.FIRMWARE
+    assert state.attributes[ATTR_INSTALLED_VERSION] == OLD_FIRMWARE
+    assert state.attributes[ATTR_LATEST_VERSION] == NEW_FIRMWARE
+
+    robot.update_firmware = AsyncMock(return_value=False)
+
+    with pytest.raises(HomeAssistantError):
+        await hass.services.async_call(
+            PLATFORM_DOMAIN,
+            SERVICE_INSTALL,
+            {ATTR_ENTITY_ID: ENTITY_ID},
+            blocking=True,
+        )
+    await hass.async_block_till_done()
+    assert robot.update_firmware.call_count == 1
+
+    robot.update_firmware = AsyncMock(return_value=True)
+    await hass.services.async_call(
+        PLATFORM_DOMAIN, SERVICE_INSTALL, {ATTR_ENTITY_ID: ENTITY_ID}, blocking=True
+    )
+    await hass.async_block_till_done()
+    assert robot.update_firmware.call_count == 1