Handle multiple files properly in zwave_js update entity (#78658)

* Handle multiple files properly in zwave_js update entity

* Until we have progress, set in progress to true. And fix when we write state

* fix tests

* Assert we set in progress to true before we get progress

* Fix tests

* Comment

homeassistant/components/zwave_js/update.py

@@ -138,7 +138,7 @@ class ZWaveNodeFirmwareUpdate(UpdateEntity):
 
     @callback
     def _unsub_firmware_events_and_reset_progress(
-        self, write_state: bool = False
+        self, write_state: bool = True
     ) -> None:
         """Unsubscribe from firmware events and reset update install progress."""
         if self._progress_unsub:
@@ -224,12 +224,14 @@ class ZWaveNodeFirmwareUpdate(UpdateEntity):
         """Install an update."""
         firmware = self._latest_version_firmware
         assert firmware
-        self._unsub_firmware_events_and_reset_progress(True)
+        self._unsub_firmware_events_and_reset_progress(False)
+        self._attr_in_progress = True
+        self.async_write_ha_state()
 
         self._progress_unsub = self.node.on(
             "firmware update progress", self._update_progress
         )
-        self._finished_unsub = self.node.once(
+        self._finished_unsub = self.node.on(
             "firmware update finished", self._update_finished
         )
 
@@ -244,6 +246,8 @@ class ZWaveNodeFirmwareUpdate(UpdateEntity):
 
             # We need to block until we receive the `firmware update finished` event
             await self._finished_event.wait()
+            # Clear the event so that a second firmware update blocks again
+            self._finished_event.clear()
             assert self._finished_status is not None
 
             # If status is not OK, we should throw an error to let the user know
@@ -262,8 +266,12 @@ class ZWaveNodeFirmwareUpdate(UpdateEntity):
             self._attr_in_progress = floor(
                 100 * self._num_files_installed / len(firmware.files)
             )
+
+            # Clear the status so we can get a new one
+            self._finished_status = None
             self.async_write_ha_state()
 
+        # If we get here, all files were installed successfully
         self._attr_installed_version = self._attr_latest_version = firmware.version
         self._latest_version_firmware = None
         self._unsub_firmware_events_and_reset_progress()
@@ -313,4 +321,4 @@ class ZWaveNodeFirmwareUpdate(UpdateEntity):
             self._poll_unsub()
             self._poll_unsub = None
 
-        self._unsub_firmware_events_and_reset_progress()
+        self._unsub_firmware_events_and_reset_progress(False)

tests/components/zwave_js/test_update.py

@@ -7,7 +7,7 @@ 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 (
+from homeassistant.components.update import (
     ATTR_AUTO_UPDATE,
     ATTR_IN_PROGRESS,
     ATTR_INSTALLED_VERSION,
@@ -54,6 +54,19 @@ FIRMWARE_UPDATES = {
     ]
 }
 
+FIRMWARE_UPDATE_MULTIPLE_FILES = {
+    "updates": [
+        {
+            "version": "11.2.4",
+            "changelog": "blah 2",
+            "files": [
+                {"target": 0, "url": "https://example2.com", "integrity": "sha2"},
+                {"target": 1, "url": "https://example4.com", "integrity": "sha4"},
+            ],
+        },
+    ]
+}
+
 
 async def test_update_entity_states(
     hass,
@@ -328,6 +341,11 @@ async def test_update_entity_progress(
     # Sleep so that task starts
     await asyncio.sleep(0.1)
 
+    state = hass.states.get(UPDATE_ENTITY)
+    assert state
+    attrs = state.attributes
+    assert attrs[ATTR_IN_PROGRESS] is True
+
     event = Event(
         type="firmware update progress",
         data={
@@ -363,7 +381,142 @@ async def test_update_entity_progress(
     state = hass.states.get(UPDATE_ENTITY)
     assert state
     attrs = state.attributes
-    assert attrs[ATTR_IN_PROGRESS] is False
+    assert attrs[ATTR_IN_PROGRESS] == 0
+    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_progress_multiple(
+    hass,
+    client,
+    climate_radio_thermostat_ct100_plus_different_endpoints,
+    integration,
+):
+    """Test update entity progress with multiple files."""
+    node = climate_radio_thermostat_ct100_plus_different_endpoints
+    client.async_send_command.return_value = FIRMWARE_UPDATE_MULTIPLE_FILES
+
+    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)
+
+    state = hass.states.get(UPDATE_ENTITY)
+    assert state
+    attrs = state.attributes
+    assert attrs[ATTR_IN_PROGRESS] is True
+
+    node.receive_event(
+        Event(
+            type="firmware update progress",
+            data={
+                "source": "node",
+                "event": "firmware update progress",
+                "nodeId": node.node_id,
+                "sentFragments": 1,
+                "totalFragments": 20,
+            },
+        )
+    )
+
+    # Block so HA can do its thing
+    await asyncio.sleep(0)
+
+    # Validate that the progress is updated (two files means progress is 50% of 5)
+    state = hass.states.get(UPDATE_ENTITY)
+    assert state
+    attrs = state.attributes
+    assert attrs[ATTR_IN_PROGRESS] == 2
+
+    node.receive_event(
+        Event(
+            type="firmware update finished",
+            data={
+                "source": "node",
+                "event": "firmware update finished",
+                "nodeId": node.node_id,
+                "status": FirmwareUpdateStatus.OK_NO_RESTART,
+            },
+        )
+    )
+
+    # Block so HA can do its thing
+    await asyncio.sleep(0)
+
+    # One file done, progress should be 50%
+    state = hass.states.get(UPDATE_ENTITY)
+    assert state
+    attrs = state.attributes
+    assert attrs[ATTR_IN_PROGRESS] == 50
+
+    node.receive_event(
+        Event(
+            type="firmware update progress",
+            data={
+                "source": "node",
+                "event": "firmware update progress",
+                "nodeId": node.node_id,
+                "sentFragments": 1,
+                "totalFragments": 20,
+            },
+        )
+    )
+
+    # Block so HA can do its thing
+    await asyncio.sleep(0)
+
+    # Validate that the progress is updated (50% + 50% of 5)
+    state = hass.states.get(UPDATE_ENTITY)
+    assert state
+    attrs = state.attributes
+    assert attrs[ATTR_IN_PROGRESS] == 52
+
+    node.receive_event(
+        Event(
+            type="firmware update finished",
+            data={
+                "source": "node",
+                "event": "firmware update finished",
+                "nodeId": node.node_id,
+                "status": FirmwareUpdateStatus.OK_NO_RESTART,
+            },
+        )
+    )
+
+    # Block so HA can do its thing
+    await asyncio.sleep(0)
+
+    # 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] == 0
     assert attrs[ATTR_INSTALLED_VERSION] == "11.2.4"
     assert attrs[ATTR_LATEST_VERSION] == "11.2.4"
     assert state.state == STATE_OFF
@@ -446,7 +599,7 @@ async def test_update_entity_install_failed(
     state = hass.states.get(UPDATE_ENTITY)
     assert state
     attrs = state.attributes
-    assert attrs[ATTR_IN_PROGRESS] is False
+    assert attrs[ATTR_IN_PROGRESS] == 0
     assert attrs[ATTR_INSTALLED_VERSION] == "10.7"
     assert attrs[ATTR_LATEST_VERSION] == "11.2.4"
     assert state.state == STATE_ON