Fix Reolink firmware updates by uploading directly (#127007)

homeassistant/components/reolink/update.py

@@ -3,11 +3,10 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from datetime import datetime
 from typing import Any
 
 from reolink_aio.exceptions import ReolinkError
-from reolink_aio.software_version import NewSoftwareVersion
+from reolink_aio.software_version import NewSoftwareVersion, SoftwareVersion
 
 from homeassistant.components.update import (
     UpdateDeviceClass,
@@ -19,7 +18,12 @@ from homeassistant.core import CALLBACK_TYPE, HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.event import async_call_later
+from homeassistant.helpers.update_coordinator import (
+    CoordinatorEntity,
+    DataUpdateCoordinator,
+)
 
+from . import DEVICE_UPDATE_INTERVAL
 from .entity import (
     ReolinkChannelCoordinatorEntity,
     ReolinkChannelEntityDescription,
@@ -28,7 +32,9 @@ from .entity import (
 )
 from .util import ReolinkConfigEntry, ReolinkData
 
+RESUME_AFTER_INSTALL = 15
 POLL_AFTER_INSTALL = 120
+POLL_PROGRESS = 2
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -86,25 +92,28 @@ async def async_setup_entry(
     async_add_entities(entities)
 
 
-class ReolinkUpdateEntity(
+class ReolinkUpdateBaseEntity(
-    ReolinkChannelCoordinatorEntity,
+    CoordinatorEntity[DataUpdateCoordinator[None]], UpdateEntity
-    UpdateEntity,
 ):
-    """Base update entity class for Reolink IP cameras."""
+    """Base update entity class for Reolink."""
 
-    entity_description: ReolinkUpdateEntityDescription
     _attr_release_url = "https://reolink.com/download-center/"
 
     def __init__(
         self,
         reolink_data: ReolinkData,
-        channel: int,
+        channel: int | None,
-        entity_description: ReolinkUpdateEntityDescription,
+        coordinator: DataUpdateCoordinator[None],
     ) -> None:
         """Initialize Reolink update entity."""
-        self.entity_description = entity_description
+        CoordinatorEntity.__init__(self, coordinator)
-        super().__init__(reolink_data, channel, reolink_data.firmware_coordinator)
+        self._channel = channel
+        self._host = reolink_data.host
         self._cancel_update: CALLBACK_TYPE | None = None
+        self._cancel_resume: CALLBACK_TYPE | None = None
+        self._cancel_progress: CALLBACK_TYPE | None = None
+        self._installing: bool = False
+        self._reolink_data = reolink_data
 
     @property
     def installed_version(self) -> str | None:
@@ -123,6 +132,16 @@ class ReolinkUpdateEntity(
 
         return new_firmware.version_string
 
+    @property
+    def in_progress(self) -> bool:
+        """Update installation progress."""
+        return self._host.api.sw_upload_progress(self._channel) < 100
+
+    @property
+    def update_percentage(self) -> int:
+        """Update installation progress."""
+        return self._host.api.sw_upload_progress(self._channel)
+
     @property
     def supported_features(self) -> UpdateEntityFeature:
         """Flag supported features."""
@@ -130,8 +149,27 @@ class ReolinkUpdateEntity(
         new_firmware = self._host.api.firmware_update_available(self._channel)
         if isinstance(new_firmware, NewSoftwareVersion):
             supported_features |= UpdateEntityFeature.RELEASE_NOTES
+            supported_features |= UpdateEntityFeature.PROGRESS
         return supported_features
 
+    @property
+    def available(self) -> bool:
+        """Return True if entity is available."""
+        if self._installing or self._cancel_update is not None:
+            return True
+        return super().available
+
+    def version_is_newer(self, latest_version: str, installed_version: str) -> bool:
+        """Return True if latest_version is newer than installed_version."""
+        try:
+            installed = SoftwareVersion(installed_version)
+            latest = SoftwareVersion(latest_version)
+        except ReolinkError:
+            # when the online update API returns a unexpected string
+            return True
+
+        return latest > installed
+
     async def async_release_notes(self) -> str | None:
         """Return the release notes."""
         new_firmware = self._host.api.firmware_update_available(self._channel)
@@ -148,6 +186,11 @@ class ReolinkUpdateEntity(
         self, version: str | None, backup: bool, **kwargs: Any
     ) -> None:
         """Install the latest firmware version."""
+        self._installing = True
+        await self._pause_update_coordinator()
+        self._cancel_progress = async_call_later(
+            self.hass, POLL_PROGRESS, self._async_update_progress
+        )
         try:
             await self._host.api.update_firmware(self._channel)
         except ReolinkError as err:
@@ -159,10 +202,38 @@ class ReolinkUpdateEntity(
             self._cancel_update = async_call_later(
                 self.hass, POLL_AFTER_INSTALL, self._async_update_future
             )
+            self._cancel_resume = async_call_later(
+                self.hass, RESUME_AFTER_INSTALL, self._resume_update_coordinator
+            )
+            self._installing = False
 
-    async def _async_update_future(self, now: datetime | None = None) -> None:
+    async def _pause_update_coordinator(self) -> None:
+        """Pause updating the states using the data update coordinator (during reboots)."""
+        self._reolink_data.device_coordinator.update_interval = None
+        self._reolink_data.device_coordinator.async_set_updated_data(None)
+
+    async def _resume_update_coordinator(self, *args) -> None:
+        """Resume updating the states using the data update coordinator (after reboots)."""
+        self._reolink_data.device_coordinator.update_interval = DEVICE_UPDATE_INTERVAL
+        try:
+            await self._reolink_data.device_coordinator.async_refresh()
+        finally:
+            self._cancel_resume = None
+
+    async def _async_update_progress(self, *args) -> None:
         """Request update."""
+        self.async_write_ha_state()
+        if self._installing:
+            self._cancel_progress = async_call_later(
+                self.hass, POLL_PROGRESS, self._async_update_progress
+            )
+
+    async def _async_update_future(self, *args) -> None:
+        """Request update."""
+        try:
             await self.async_update()
+        finally:
+            self._cancel_update = None
 
     async def async_added_to_hass(self) -> None:
         """Entity created."""
@@ -176,16 +247,44 @@ class ReolinkUpdateEntity(
             self._host.firmware_ch_list.remove(self._channel)
         if self._cancel_update is not None:
             self._cancel_update()
+        if self._cancel_progress is not None:
+            self._cancel_progress()
+        if self._cancel_resume is not None:
+            self._cancel_resume()
+
+
+class ReolinkUpdateEntity(
+    ReolinkUpdateBaseEntity,
+    ReolinkChannelCoordinatorEntity,
+):
+    """Base update entity class for Reolink IP cameras."""
+
+    entity_description: ReolinkUpdateEntityDescription
+    _channel: int
+
+    def __init__(
+        self,
+        reolink_data: ReolinkData,
+        channel: int,
+        entity_description: ReolinkUpdateEntityDescription,
+    ) -> None:
+        """Initialize Reolink update entity."""
+        self.entity_description = entity_description
+        ReolinkUpdateBaseEntity.__init__(
+            self, reolink_data, channel, reolink_data.firmware_coordinator
+        )
+        ReolinkChannelCoordinatorEntity.__init__(
+            self, reolink_data, channel, reolink_data.firmware_coordinator
+        )
 
 
 class ReolinkHostUpdateEntity(
+    ReolinkUpdateBaseEntity,
     ReolinkHostCoordinatorEntity,
-    UpdateEntity,
 ):
     """Update entity class for Reolink Host."""
 
     entity_description: ReolinkHostUpdateEntityDescription
-    _attr_release_url = "https://reolink.com/download-center/"
 
     def __init__(
         self,
@@ -194,76 +293,9 @@ class ReolinkHostUpdateEntity(
     ) -> None:
         """Initialize Reolink update entity."""
         self.entity_description = entity_description
-        super().__init__(reolink_data, reolink_data.firmware_coordinator)
+        ReolinkUpdateBaseEntity.__init__(
-        self._cancel_update: CALLBACK_TYPE | None = None
+            self, reolink_data, None, reolink_data.firmware_coordinator
-
-    @property
-    def installed_version(self) -> str | None:
-        """Version currently in use."""
-        return self._host.api.sw_version
-
-    @property
-    def latest_version(self) -> str | None:
-        """Latest version available for install."""
-        new_firmware = self._host.api.firmware_update_available()
-        if not new_firmware:
-            return self.installed_version
-
-        if isinstance(new_firmware, str):
-            return new_firmware
-
-        return new_firmware.version_string
-
-    @property
-    def supported_features(self) -> UpdateEntityFeature:
-        """Flag supported features."""
-        supported_features = UpdateEntityFeature.INSTALL
-        new_firmware = self._host.api.firmware_update_available()
-        if isinstance(new_firmware, NewSoftwareVersion):
-            supported_features |= UpdateEntityFeature.RELEASE_NOTES
-        return supported_features
-
-    async def async_release_notes(self) -> str | None:
-        """Return the release notes."""
-        new_firmware = self._host.api.firmware_update_available()
-        assert isinstance(new_firmware, NewSoftwareVersion)
-
-        return (
-            "If the install button fails, download this"
-            f" [firmware zip file]({new_firmware.download_url})."
-            " Then, follow the installation guide (PDF in the zip file).\n\n"
-            f"## Release notes\n\n{new_firmware.release_notes}"
         )
-
+        ReolinkHostCoordinatorEntity.__init__(
-    async def async_install(
+            self, reolink_data, reolink_data.firmware_coordinator
-        self, version: str | None, backup: bool, **kwargs: Any
-    ) -> None:
-        """Install the latest firmware version."""
-        try:
-            await self._host.api.update_firmware()
-        except ReolinkError as err:
-            raise HomeAssistantError(
-                f"Error trying to update Reolink firmware: {err}"
-            ) from err
-        finally:
-            self.async_write_ha_state()
-            self._cancel_update = async_call_later(
-                self.hass, POLL_AFTER_INSTALL, self._async_update_future
         )
-
-    async def _async_update_future(self, now: datetime | None = None) -> None:
-        """Request update."""
-        await self.async_update()
-
-    async def async_added_to_hass(self) -> None:
-        """Entity created."""
-        await super().async_added_to_hass()
-        self._host.firmware_ch_list.append(None)
-
-    async def async_will_remove_from_hass(self) -> None:
-        """Entity removed."""
-        await super().async_will_remove_from_hass()
-        if None in self._host.firmware_ch_list:
-            self._host.firmware_ch_list.remove(None)
-        if self._cancel_update is not None:
-            self._cancel_update()

tests/components/reolink/conftest.py

@@ -86,6 +86,7 @@ def reolink_connect_class() -> Generator[MagicMock]:
         host_mock.sw_version_update_required = False
         host_mock.hardware_version = "IPC_00000"
         host_mock.sw_version = "v1.0.0.0.0.0000"
+        host_mock.sw_upload_progress.return_value = 100
         host_mock.manufacturer = "Reolink"
         host_mock.model = TEST_HOST_MODEL
         host_mock.item_number = TEST_ITEM_NUMBER

tests/components/reolink/test_update.py

@@ -1,5 +1,7 @@
 """Test the Reolink update platform."""
 
+import asyncio
+from datetime import timedelta
 from unittest.mock import MagicMock, patch
 
 from freezegun.api import FrozenDateTimeFactory
@@ -7,12 +9,13 @@ import pytest
 from reolink_aio.exceptions import ReolinkError
 from reolink_aio.software_version import NewSoftwareVersion
 
-from homeassistant.components.reolink.update import POLL_AFTER_INSTALL
+from homeassistant.components.reolink.update import POLL_AFTER_INSTALL, POLL_PROGRESS
 from homeassistant.components.update import DOMAIN as UPDATE_DOMAIN, SERVICE_INSTALL
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.const import ATTR_ENTITY_ID, STATE_OFF, STATE_ON, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
+from homeassistant.util.dt import utcnow
 
 from .conftest import TEST_CAM_NAME, TEST_NVR_NAME
 
@@ -73,6 +76,7 @@ async def test_update_firm(
 ) -> None:
     """Test update state when update available with firmware info from reolink.com."""
     reolink_connect.camera_name.return_value = TEST_CAM_NAME
+    reolink_connect.sw_upload_progress.return_value = 100
     reolink_connect.camera_sw_version.return_value = "v1.1.0.0.0.0000"
     new_firmware = NewSoftwareVersion(
         version_string="v3.3.0.226_23031644",
@@ -88,6 +92,8 @@ async def test_update_firm(
 
     entity_id = f"{Platform.UPDATE}.{entity_name}_firmware"
     assert hass.states.get(entity_id).state == STATE_ON
+    assert not hass.states.get(entity_id).attributes["in_progress"]
+    assert hass.states.get(entity_id).attributes["update_percentage"] is None
 
     # release notes
     client = await hass_ws_client(hass)
@@ -113,6 +119,22 @@ async def test_update_firm(
     )
     reolink_connect.update_firmware.assert_called()
 
+    reolink_connect.sw_upload_progress.return_value = 50
+    freezer.tick(POLL_PROGRESS)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+
+    assert hass.states.get(entity_id).attributes["in_progress"]
+    assert hass.states.get(entity_id).attributes["update_percentage"] == 50
+
+    reolink_connect.sw_upload_progress.return_value = 100
+    freezer.tick(POLL_AFTER_INSTALL)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+
+    assert not hass.states.get(entity_id).attributes["in_progress"]
+    assert hass.states.get(entity_id).attributes["update_percentage"] is None
+
     reolink_connect.update_firmware.side_effect = ReolinkError("Test error")
     with pytest.raises(HomeAssistantError):
         await hass.services.async_call(
@@ -132,3 +154,53 @@ async def test_update_firm(
     assert hass.states.get(entity_id).state == STATE_OFF
 
     reolink_connect.update_firmware.side_effect = None
+
+
+@pytest.mark.parametrize("entity_name", [TEST_NVR_NAME, TEST_CAM_NAME])
+async def test_update_firm_keeps_available(
+    hass: HomeAssistant,
+    config_entry: MockConfigEntry,
+    reolink_connect: MagicMock,
+    hass_ws_client: WebSocketGenerator,
+    entity_name: str,
+) -> None:
+    """Test update entity keeps being available during update."""
+    reolink_connect.camera_name.return_value = TEST_CAM_NAME
+    reolink_connect.camera_sw_version.return_value = "v1.1.0.0.0.0000"
+    new_firmware = NewSoftwareVersion(
+        version_string="v3.3.0.226_23031644",
+        download_url=TEST_DOWNLOAD_URL,
+        release_notes=TEST_RELEASE_NOTES,
+    )
+    reolink_connect.firmware_update_available.return_value = new_firmware
+
+    with patch("homeassistant.components.reolink.PLATFORMS", [Platform.UPDATE]):
+        assert await hass.config_entries.async_setup(config_entry.entry_id)
+    await hass.async_block_till_done()
+    assert config_entry.state is ConfigEntryState.LOADED
+
+    entity_id = f"{Platform.UPDATE}.{entity_name}_firmware"
+    assert hass.states.get(entity_id).state == STATE_ON
+
+    async def mock_update_firmware(*args, **kwargs) -> None:
+        await asyncio.sleep(0.000005)
+
+    reolink_connect.update_firmware = mock_update_firmware
+
+    # test install
+    with patch("homeassistant.components.reolink.update.POLL_PROGRESS", 0.000001):
+        await hass.services.async_call(
+            UPDATE_DOMAIN,
+            SERVICE_INSTALL,
+            {ATTR_ENTITY_ID: entity_id},
+            blocking=True,
+        )
+
+    reolink_connect.session_active = False
+    async_fire_time_changed(hass, utcnow() + timedelta(seconds=1))
+    await hass.async_block_till_done()
+
+    # still available
+    assert hass.states.get(entity_id).state == STATE_ON
+
+    reolink_connect.session_active = True