Reset update percentage after install

tests/test_update.py

@@ -1,5 +1,6 @@
 """Test ZHA firmware updates."""
 
+import asyncio
 from unittest.mock import ANY, AsyncMock, call, patch
 
 import pytest
@@ -413,6 +414,8 @@ def read_new_fw_version(*args, **kwargs):
         == f"0x{fw_image.firmware.header.file_version:08x}"
     )
     assert not entity.state[ATTR_IN_PROGRESS]
+    # The progress percentage is cleared once the update finishes
+    assert entity.state[ATTR_UPDATE_PERCENTAGE] is None
     assert entity.state[ATTR_LATEST_VERSION] == entity.state[ATTR_INSTALLED_VERSION]
 
     # If we send a progress notification incorrectly, it won't be handled
@@ -505,6 +508,10 @@ async def endpoint_reply(cluster, sequence, data, **kwargs):
         )
         await zha_gateway.async_block_till_done()
 
+    # The progress state is cleared even when the update fails
+    assert not entity.state[ATTR_IN_PROGRESS]
+    assert entity.state[ATTR_UPDATE_PERCENTAGE] is None
+
 
 async def test_firmware_update_empty_exception_message(zha_gateway: Gateway) -> None:
     """Bare ``TimeoutError()`` from zigpy must still produce an identifiable message."""
@@ -549,6 +556,47 @@ async def test_firmware_update_empty_exception_message(zha_gateway: Gateway) ->
     assert not entity.state[ATTR_IN_PROGRESS]
 
 
+async def test_firmware_update_cancelled(zha_gateway: Gateway) -> None:
+    """A cancelled install must propagate and still clear the progress state."""
+    zigpy_device = zigpy_device_mock(zha_gateway)
+    zha_device, ota_cluster, fw_image, installed_fw_version = await setup_test_data(
+        zha_gateway, zigpy_device
+    )
+
+    entity = get_entity(zha_device, platform=Platform.UPDATE)
+
+    await ota_cluster._handle_query_next_image(
+        foundation.ZCLHeader.cluster(
+            tsn=0x12, command_id=general.Ota.ServerCommandDefs.query_next_image.id
+        ),
+        general.QueryNextImageCommand(
+            field_control=fw_image.firmware.header.field_control,
+            manufacturer_code=zha_device.manufacturer_code,
+            image_type=fw_image.firmware.header.image_type,
+            current_file_version=installed_fw_version,
+            hardware_version=1,
+        ),
+    )
+    await zha_gateway.async_block_till_done()
+
+    # `CancelledError` is a `BaseException`, so it must propagate untouched
+    # rather than being wrapped in a `ZHAException`.
+    with (
+        patch(
+            "zigpy.device.Device.update_firmware",
+            AsyncMock(side_effect=asyncio.CancelledError),
+        ),
+        pytest.raises(asyncio.CancelledError),
+    ):
+        await entity.async_install(
+            version=f"0x{fw_image.firmware.header.file_version:08x}"
+        )
+
+    # The progress state is still cleared when the install is cancelled
+    assert not entity.state[ATTR_IN_PROGRESS]
+    assert entity.state[ATTR_UPDATE_PERCENTAGE] is None
+
+
 async def test_firmware_update_downgrade(zha_gateway: Gateway) -> None:
     """Test ZHA update platform - force a firmware downgrade."""
     zigpy_device = zigpy_device_mock(zha_gateway)

zha/application/platforms/update.py

@@ -265,22 +265,19 @@ async def async_install(self, version: str | None) -> None:
                 progress_callback=self._update_progress,
             )
         except Exception as ex:
-            self._attr_in_progress = False
-            self.maybe_emit_state_changed_event()
             raise ZHAException(
                 f"Update was not successful: {str(ex) or repr(ex)}"
             ) from ex
+        finally:
+            # Clear the progress state, regardless of the update's outcome
+            self._attr_in_progress = False
+            self._attr_update_percentage = None
+            self.maybe_emit_state_changed_event()
 
         # If the update finished but was not successful, we should also throw an error
         if result != Status.SUCCESS:
-            self._attr_in_progress = False
-            self.maybe_emit_state_changed_event()
             raise ZHAException(f"Update was not successful: {result}")
 
-        # Clear the state
-        self._attr_in_progress = False
-        self.maybe_emit_state_changed_event()
-
     async def on_remove(self) -> None:
         """Call when entity will be removed."""
         self._attr_in_progress = False