Add more features to lamarzocco updates (#143157)

homeassistant/components/lamarzocco/update.py

@@ -1,9 +1,10 @@
 """Support for La Marzocco update entities."""
 
+import asyncio
 from dataclasses import dataclass
 from typing import Any
 
-from pylamarzocco.const import FirmwareType
+from pylamarzocco.const import FirmwareType, UpdateCommandStatus
 from pylamarzocco.exceptions import RequestNotSuccessful
 
 from homeassistant.components.update import (
@@ -22,6 +23,7 @@ from .coordinator import LaMarzoccoConfigEntry
 from .entity import LaMarzoccoEntity, LaMarzoccoEntityDescription
 
 PARALLEL_UPDATES = 1
+MAX_UPDATE_WAIT = 150
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -71,7 +73,11 @@ class LaMarzoccoUpdateEntity(LaMarzoccoEntity, UpdateEntity):
     """Entity representing the update state."""
 
     entity_description: LaMarzoccoUpdateEntityDescription
-    _attr_supported_features = UpdateEntityFeature.INSTALL
+    _attr_supported_features = (
+        UpdateEntityFeature.INSTALL
+        | UpdateEntityFeature.PROGRESS
+        | UpdateEntityFeature.RELEASE_NOTES
+    )
 
     @property
     def installed_version(self) -> str:
@@ -94,15 +100,40 @@ class LaMarzoccoUpdateEntity(LaMarzoccoEntity, UpdateEntity):
         """Return the release notes URL."""
         return "https://support-iot.lamarzocco.com/firmware-updates/"
 
+    def release_notes(self) -> str | None:
+        """Return the release notes for the latest firmware version."""
+        if available_update := self.coordinator.device.settings.firmwares[
+            self.entity_description.component
+        ].available_update:
+            return available_update.change_log
+        return None
+
     async def async_install(
         self, version: str | None, backup: bool, **kwargs: Any
     ) -> None:
         """Install an update."""
+
         self._attr_in_progress = True
         self.async_write_ha_state()
+
+        counter = 0
+
+        def _raise_timeout_error() -> None:  # to avoid TRY301
+            raise TimeoutError("Update timed out")
+
         try:
             await self.coordinator.device.update_firmware()
-        except RequestNotSuccessful as exc:
+            while (
+                update_progress := await self.coordinator.device.get_firmware()
+            ).command_status is UpdateCommandStatus.IN_PROGRESS:
+                if counter >= MAX_UPDATE_WAIT:
+                    _raise_timeout_error()
+                self._attr_update_percentage = update_progress.progress_percentage
+                self.async_write_ha_state()
+                await asyncio.sleep(3)
+                counter += 1
+
+        except (TimeoutError, RequestNotSuccessful) as exc:
             raise HomeAssistantError(
                 translation_domain=DOMAIN,
                 translation_key="update_failed",
@@ -110,5 +141,6 @@ class LaMarzoccoUpdateEntity(LaMarzoccoEntity, UpdateEntity):
                     "key": self.entity_description.key,
                 },
             ) from exc
-        self._attr_in_progress = False
-        await self.coordinator.async_request_refresh()
+        finally:
+            self._attr_in_progress = False
+            await self.coordinator.async_request_refresh()

tests/components/lamarzocco/snapshots/test_update.ambr

@@ -27,7 +27,7 @@
     'original_name': 'Gateway firmware',
     'platform': 'lamarzocco',
     'previous_unique_id': None,
-    'supported_features': <UpdateEntityFeature: 1>,
+    'supported_features': <UpdateEntityFeature: 21>,
     'translation_key': 'gateway_firmware',
     'unique_id': 'GS012345_gateway_firmware',
     'unit_of_measurement': None,
@@ -47,7 +47,7 @@
       'release_summary': None,
       'release_url': 'https://support-iot.lamarzocco.com/firmware-updates/',
       'skipped_version': None,
-      'supported_features': <UpdateEntityFeature: 1>,
+      'supported_features': <UpdateEntityFeature: 21>,
       'title': None,
       'update_percentage': None,
     }),
@@ -87,7 +87,7 @@
     'original_name': 'Machine firmware',
     'platform': 'lamarzocco',
     'previous_unique_id': None,
-    'supported_features': <UpdateEntityFeature: 1>,
+    'supported_features': <UpdateEntityFeature: 21>,
     'translation_key': 'machine_firmware',
     'unique_id': 'GS012345_machine_firmware',
     'unit_of_measurement': None,
@@ -107,7 +107,7 @@
       'release_summary': None,
       'release_url': 'https://support-iot.lamarzocco.com/firmware-updates/',
       'skipped_version': None,
-      'supported_features': <UpdateEntityFeature: 1>,
+      'supported_features': <UpdateEntityFeature: 21>,
       'title': None,
       'update_percentage': None,
     }),

tests/components/lamarzocco/test_update.py

@@ -1,8 +1,16 @@
 """Tests for the La Marzocco Update Entities."""
 
-from unittest.mock import MagicMock, patch
+from collections.abc import Generator
+from unittest.mock import AsyncMock, MagicMock, patch
 
+from pylamarzocco.const import (
+    FirmwareType,
+    UpdateCommandStatus,
+    UpdateProgressInfo,
+    UpdateStatus,
+)
 from pylamarzocco.exceptions import RequestNotSuccessful
+from pylamarzocco.models import UpdateDetails
 import pytest
 from syrupy import SnapshotAssertion
 
@@ -15,6 +23,17 @@ from homeassistant.helpers import entity_registry as er
 from . import async_init_integration
 
 from tests.common import MockConfigEntry, snapshot_platform
+from tests.typing import WebSocketGenerator
+
+
+@pytest.fixture(autouse=True)
+def mock_sleep() -> Generator[AsyncMock]:
+    """Mock asyncio.sleep."""
+    with patch(
+        "homeassistant.components.lamarzocco.update.asyncio.sleep",
+        return_value=AsyncMock(),
+    ) as mock_sleep:
+        yield mock_sleep
 
 
 async def test_update(
@@ -29,17 +48,51 @@ async def test_update(
     await snapshot_platform(hass, entity_registry, snapshot, mock_config_entry.entry_id)
 
 
-async def test_update_entites(
+async def test_update_process(
     hass: HomeAssistant,
     mock_lamarzocco: MagicMock,
     mock_config_entry: MockConfigEntry,
+    hass_ws_client: WebSocketGenerator,
 ) -> None:
     """Test the La Marzocco update entities."""
 
     serial_number = mock_lamarzocco.serial_number
 
+    mock_lamarzocco.get_firmware.side_effect = [
+        UpdateDetails(
+            status=UpdateStatus.TO_UPDATE,
+            command_status=UpdateCommandStatus.IN_PROGRESS,
+            progress_info=UpdateProgressInfo.STARTING_PROCESS,
+            progress_percentage=0,
+        ),
+        UpdateDetails(
+            status=UpdateStatus.UPDATED,
+            command_status=None,
+            progress_info=None,
+            progress_percentage=None,
+        ),
+    ]
+
     await async_init_integration(hass, mock_config_entry)
 
+    client = await hass_ws_client(hass)
+    await hass.async_block_till_done()
+
+    await client.send_json(
+        {
+            "id": 1,
+            "type": "update/release_notes",
+            "entity_id": f"update.{serial_number}_gateway_firmware",
+        }
+    )
+    result = await client.receive_json()
+    assert (
+        mock_lamarzocco.settings.firmwares[
+            FirmwareType.GATEWAY
+        ].available_update.change_log
+        in result["result"]
+    )
+
     await hass.services.async_call(
         UPDATE_DOMAIN,
         SERVICE_INSTALL,
@@ -76,3 +129,35 @@ async def test_update_error(
             blocking=True,
         )
     assert exc_info.value.translation_key == "update_failed"
+
+
+async def test_update_times_out(
+    hass: HomeAssistant,
+    mock_lamarzocco: MagicMock,
+    mock_config_entry: MockConfigEntry,
+) -> None:
+    """Test error during update."""
+    mock_lamarzocco.get_firmware.return_value = UpdateDetails(
+        status=UpdateStatus.TO_UPDATE,
+        command_status=UpdateCommandStatus.IN_PROGRESS,
+        progress_info=UpdateProgressInfo.STARTING_PROCESS,
+        progress_percentage=0,
+    )
+    await async_init_integration(hass, mock_config_entry)
+
+    state = hass.states.get(f"update.{mock_lamarzocco.serial_number}_gateway_firmware")
+    assert state
+
+    with (
+        patch("homeassistant.components.lamarzocco.update.MAX_UPDATE_WAIT", 0),
+        pytest.raises(HomeAssistantError) as exc_info,
+    ):
+        await hass.services.async_call(
+            UPDATE_DOMAIN,
+            SERVICE_INSTALL,
+            {
+                ATTR_ENTITY_ID: f"update.{mock_lamarzocco.serial_number}_gateway_firmware",
+            },
+            blocking=True,
+        )
+    assert exc_info.value.translation_key == "update_failed"