diff --git a/tests/test_update.py b/tests/test_update.py index 987bd535c..cd3b82b23 100644 --- a/tests/test_update.py +++ b/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) diff --git a/zha/application/platforms/update.py b/zha/application/platforms/update.py index cc7ce0a92..2ef177757 100644 --- a/zha/application/platforms/update.py +++ b/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