diff --git a/homeassistant/components/zwave_js/update.py b/homeassistant/components/zwave_js/update.py index 4f25d138aea..932ed46a0fc 100644 --- a/homeassistant/components/zwave_js/update.py +++ b/homeassistant/components/zwave_js/update.py @@ -12,7 +12,12 @@ from zwave_js_server.client import Client as ZwaveClient from zwave_js_server.const import NodeStatus from zwave_js_server.exceptions import BaseZwaveJSServerError, FailedZWaveCommand from zwave_js_server.model.driver import Driver -from zwave_js_server.model.firmware import FirmwareUpdateInfo, FirmwareUpdateProgress +from zwave_js_server.model.firmware import ( + FirmwareUpdateFinished, + FirmwareUpdateInfo, + FirmwareUpdateProgress, + FirmwareUpdateStatus, +) from zwave_js_server.model.node import Node as ZwaveNode from homeassistant.components.update import UpdateDeviceClass, UpdateEntity @@ -82,7 +87,10 @@ class ZWaveNodeFirmwareUpdate(UpdateEntity): self._status_unsub: Callable[[], None] | None = None self._poll_unsub: Callable[[], None] | None = None self._progress_unsub: Callable[[], None] | None = None + self._finished_unsub: Callable[[], None] | None = None self._num_files_installed: int = 0 + self._finished_event = asyncio.Event() + self._finished_status: FirmwareUpdateStatus | None = None # Entity class attributes self._attr_name = "Firmware" @@ -119,18 +127,38 @@ class ZWaveNodeFirmwareUpdate(UpdateEntity): self.async_write_ha_state() @callback - def _reset_progress(self) -> None: - """Reset update install progress.""" + def _update_finished(self, event: dict[str, Any]) -> None: + """Update install progress on event.""" + finished: FirmwareUpdateFinished = event["firmware_update_finished"] + self._finished_status = finished.status + self._finished_event.set() + + @callback + def _unsub_firmware_events_and_reset_progress( + self, write_state: bool = False + ) -> None: + """Unsubscribe from firmware events and reset update install progress.""" if self._progress_unsub: self._progress_unsub() self._progress_unsub = None + + if self._finished_unsub: + self._finished_unsub() + self._finished_unsub = None + + self._finished_status = None + self._finished_event.clear() self._num_files_installed = 0 - self._attr_in_progress = False - self.async_write_ha_state() + self._attr_in_progress = 0 + if write_state: + self.async_write_ha_state() async def _async_update(self, _: HomeAssistant | datetime | None = None) -> None: """Update the entity.""" self._poll_unsub = None + + # If device is asleep/dead, wait for it to wake up/become alive before + # attempting an update for status, event_name in ( (NodeStatus.ASLEEP, "wake up"), (NodeStatus.DEAD, "alive"), @@ -187,19 +215,40 @@ class ZWaveNodeFirmwareUpdate(UpdateEntity): """Install an update.""" firmware = self._latest_version_firmware assert firmware - self._attr_in_progress = 0 - self.async_write_ha_state() + self._unsub_firmware_events_and_reset_progress(True) + self._progress_unsub = self.node.on( "firmware update progress", self._update_progress ) + self._finished_unsub = self.node.once( + "firmware update finished", self._update_finished + ) + for file in firmware.files: try: await self.driver.controller.async_begin_ota_firmware_update( self.node, file ) except BaseZwaveJSServerError as err: - self._reset_progress() + self._unsub_firmware_events_and_reset_progress() raise HomeAssistantError(err) from err + + # We need to block until we receive the `firmware update finished` event + await self._finished_event.wait() + assert self._finished_status is not None + + # If status is not OK, we should throw an error to let the user know + if self._finished_status not in ( + FirmwareUpdateStatus.OK_NO_RESTART, + FirmwareUpdateStatus.OK_RESTART_PENDING, + FirmwareUpdateStatus.OK_WAITING_FOR_ACTIVATION, + ): + status = self._finished_status + self._unsub_firmware_events_and_reset_progress() + raise HomeAssistantError(status.name.replace("_", " ").title()) + + # If we get here, the firmware installation was successful and we need to + # update progress accordingly self._num_files_installed += 1 self._attr_in_progress = floor( 100 * self._num_files_installed / len(firmware.files) @@ -208,7 +257,7 @@ class ZWaveNodeFirmwareUpdate(UpdateEntity): self._attr_installed_version = self._attr_latest_version = firmware.version self._latest_version_firmware = None - self._reset_progress() + self._unsub_firmware_events_and_reset_progress() async def async_poll_value(self, _: bool) -> None: """Poll a value.""" @@ -255,6 +304,4 @@ class ZWaveNodeFirmwareUpdate(UpdateEntity): self._poll_unsub() self._poll_unsub = None - if self._progress_unsub: - self._progress_unsub() - self._progress_unsub = None + self._unsub_firmware_events_and_reset_progress() diff --git a/tests/components/zwave_js/test_update.py b/tests/components/zwave_js/test_update.py index 76fecfdee6d..0b567d93106 100644 --- a/tests/components/zwave_js/test_update.py +++ b/tests/components/zwave_js/test_update.py @@ -1,9 +1,11 @@ """Test the Z-Wave JS update entities.""" +import asyncio from datetime import timedelta import pytest from zwave_js_server.event import Event from zwave_js_server.exceptions import FailedZWaveCommand +from zwave_js_server.model.firmware import FirmwareUpdateStatus from homeassistant.components.update.const import ( ATTR_AUTO_UPDATE, @@ -51,7 +53,7 @@ FIRMWARE_UPDATES = { } -async def test_update_entity_success( +async def test_update_entity_states( hass, client, climate_radio_thermostat_ct100_plus_different_endpoints, @@ -60,7 +62,7 @@ async def test_update_entity_success( caplog, hass_ws_client, ): - """Test update entity.""" + """Test update entity states.""" ws_client = await hass_ws_client(hass) await hass.async_block_till_done() @@ -137,39 +139,14 @@ async def test_update_entity_success( client.async_send_command.reset_mock() - # Test successful install call without a version - await hass.services.async_call( - UPDATE_DOMAIN, - SERVICE_INSTALL, - { - ATTR_ENTITY_ID: UPDATE_ENTITY, - }, - blocking=True, - ) - args = client.async_send_command.call_args_list[0][0][0] - assert args["command"] == "controller.begin_ota_firmware_update" - assert ( - args["nodeId"] - == climate_radio_thermostat_ct100_plus_different_endpoints.node_id - ) - assert args["update"] == { - "target": 0, - "url": "https://example2.com", - "integrity": "sha2", - } - - client.async_send_command.reset_mock() - - -async def test_update_entity_install_failure( +async def test_update_entity_install_raises( hass, client, climate_radio_thermostat_ct100_plus_different_endpoints, - controller_node, integration, ): - """Test update entity failed install.""" + """Test update entity install raises exception.""" client.async_send_command.return_value = FIRMWARE_UPDATES async_fire_time_changed(hass, dt_util.utcnow() + timedelta(days=1)) @@ -287,11 +264,10 @@ async def test_update_entity_ha_not_running( assert args["nodeId"] == zen_31.node_id -async def test_update_entity_failure( +async def test_update_entity_update_failure( hass, client, climate_radio_thermostat_ct100_plus_different_endpoints, - controller_node, integration, ): """Test update entity update failed.""" @@ -311,3 +287,169 @@ async def test_update_entity_failure( args["nodeId"] == climate_radio_thermostat_ct100_plus_different_endpoints.node_id ) + + +async def test_update_entity_progress( + hass, + client, + climate_radio_thermostat_ct100_plus_different_endpoints, + integration, +): + """Test update entity progress.""" + node = climate_radio_thermostat_ct100_plus_different_endpoints + client.async_send_command.return_value = FIRMWARE_UPDATES + + async_fire_time_changed(hass, dt_util.utcnow() + timedelta(days=1)) + await hass.async_block_till_done() + + state = hass.states.get(UPDATE_ENTITY) + assert state + assert state.state == STATE_ON + attrs = state.attributes + assert attrs[ATTR_INSTALLED_VERSION] == "10.7" + assert attrs[ATTR_LATEST_VERSION] == "11.2.4" + + client.async_send_command.reset_mock() + client.async_send_command.return_value = None + + # Test successful install call without a version + install_task = hass.async_create_task( + hass.services.async_call( + UPDATE_DOMAIN, + SERVICE_INSTALL, + { + ATTR_ENTITY_ID: UPDATE_ENTITY, + }, + blocking=True, + ) + ) + + # Sleep so that task starts + await asyncio.sleep(0.1) + + event = Event( + type="firmware update progress", + data={ + "source": "node", + "event": "firmware update progress", + "nodeId": node.node_id, + "sentFragments": 1, + "totalFragments": 20, + }, + ) + node.receive_event(event) + + # Validate that the progress is updated + state = hass.states.get(UPDATE_ENTITY) + assert state + attrs = state.attributes + assert attrs[ATTR_IN_PROGRESS] == 5 + + event = Event( + type="firmware update finished", + data={ + "source": "node", + "event": "firmware update finished", + "nodeId": node.node_id, + "status": FirmwareUpdateStatus.OK_NO_RESTART, + }, + ) + + node.receive_event(event) + await hass.async_block_till_done() + + # Validate that progress is reset and entity reflects new version + state = hass.states.get(UPDATE_ENTITY) + assert state + attrs = state.attributes + assert attrs[ATTR_IN_PROGRESS] is False + assert attrs[ATTR_INSTALLED_VERSION] == "11.2.4" + assert attrs[ATTR_LATEST_VERSION] == "11.2.4" + assert state.state == STATE_OFF + + await install_task + + +async def test_update_entity_install_failed( + hass, + client, + climate_radio_thermostat_ct100_plus_different_endpoints, + integration, + caplog, +): + """Test update entity install returns error status.""" + node = climate_radio_thermostat_ct100_plus_different_endpoints + client.async_send_command.return_value = FIRMWARE_UPDATES + + async_fire_time_changed(hass, dt_util.utcnow() + timedelta(days=1)) + await hass.async_block_till_done() + + state = hass.states.get(UPDATE_ENTITY) + assert state + assert state.state == STATE_ON + attrs = state.attributes + assert attrs[ATTR_INSTALLED_VERSION] == "10.7" + assert attrs[ATTR_LATEST_VERSION] == "11.2.4" + + client.async_send_command.reset_mock() + client.async_send_command.return_value = None + + async def call_install(): + await hass.services.async_call( + UPDATE_DOMAIN, + SERVICE_INSTALL, + { + ATTR_ENTITY_ID: UPDATE_ENTITY, + }, + blocking=True, + ) + + # Test install call - we expect it to raise + install_task = hass.async_create_task(call_install()) + + # Sleep so that task starts + await asyncio.sleep(0.1) + + event = Event( + type="firmware update progress", + data={ + "source": "node", + "event": "firmware update progress", + "nodeId": node.node_id, + "sentFragments": 1, + "totalFragments": 20, + }, + ) + node.receive_event(event) + + # Validate that the progress is updated + state = hass.states.get(UPDATE_ENTITY) + assert state + attrs = state.attributes + assert attrs[ATTR_IN_PROGRESS] == 5 + + event = Event( + type="firmware update finished", + data={ + "source": "node", + "event": "firmware update finished", + "nodeId": node.node_id, + "status": FirmwareUpdateStatus.ERROR_TIMEOUT, + }, + ) + + node.receive_event(event) + await hass.async_block_till_done() + + # Validate that progress is reset and entity reflects old version + state = hass.states.get(UPDATE_ENTITY) + assert state + attrs = state.attributes + assert attrs[ATTR_IN_PROGRESS] is False + assert attrs[ATTR_INSTALLED_VERSION] == "10.7" + assert attrs[ATTR_LATEST_VERSION] == "11.2.4" + assert state.state == STATE_ON + + # validate that the install task failed + with pytest.raises(HomeAssistantError): + await install_task