From 52d2ebd2c8af4aaea95bc1310c658f9223a94fd4 Mon Sep 17 00:00:00 2001
From: Raman Gupta <7243222+raman325@users.noreply.github.com>
Date: Thu, 8 Sep 2022 16:40:55 -0400
Subject: [PATCH] Show progress for zwave_js.update entity (#77905)

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>
---
 homeassistant/components/zwave_js/update.py | 68 ++++++++++++++++++---
 1 file changed, 59 insertions(+), 9 deletions(-)

diff --git a/homeassistant/components/zwave_js/update.py b/homeassistant/components/zwave_js/update.py
index d106c7d6dd3..4f25d138aea 100644
--- a/homeassistant/components/zwave_js/update.py
+++ b/homeassistant/components/zwave_js/update.py
@@ -4,6 +4,7 @@ from __future__ import annotations
 import asyncio
 from collections.abc import Callable
 from datetime import datetime, timedelta
+from math import floor
 from typing import Any
 
 from awesomeversion import AwesomeVersion
@@ -11,7 +12,7 @@ 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
+from zwave_js_server.model.firmware import FirmwareUpdateInfo, FirmwareUpdateProgress
 from zwave_js_server.model.node import Node as ZwaveNode
 
 from homeassistant.components.update import UpdateDeviceClass, UpdateEntity
@@ -63,7 +64,9 @@ class ZWaveNodeFirmwareUpdate(UpdateEntity):
     _attr_entity_category = EntityCategory.CONFIG
     _attr_device_class = UpdateDeviceClass.FIRMWARE
     _attr_supported_features = (
-        UpdateEntityFeature.INSTALL | UpdateEntityFeature.RELEASE_NOTES
+        UpdateEntityFeature.INSTALL
+        | UpdateEntityFeature.RELEASE_NOTES
+        | UpdateEntityFeature.PROGRESS
     )
     _attr_has_entity_name = True
     _attr_should_poll = False
@@ -78,6 +81,8 @@ class ZWaveNodeFirmwareUpdate(UpdateEntity):
         self._latest_version_firmware: FirmwareUpdateInfo | None = None
         self._status_unsub: Callable[[], None] | None = None
         self._poll_unsub: Callable[[], None] | None = None
+        self._progress_unsub: Callable[[], None] | None = None
+        self._num_files_installed: int = 0
 
         # Entity class attributes
         self._attr_name = "Firmware"
@@ -93,6 +98,36 @@ class ZWaveNodeFirmwareUpdate(UpdateEntity):
         self._status_unsub = None
         self.hass.async_create_task(self._async_update())
 
+    @callback
+    def _update_progress(self, event: dict[str, Any]) -> None:
+        """Update install progress on event."""
+        progress: FirmwareUpdateProgress = event["firmware_update_progress"]
+        if not self._latest_version_firmware:
+            return
+        # We will assume that each file in the firmware update represents an equal
+        # percentage of the overall progress. This is likely not true because each file
+        # may be a different size, but it's the best we can do since we don't know the
+        # total number of fragments across all files.
+        self._attr_in_progress = floor(
+            100
+            * (
+                self._num_files_installed
+                + (progress.sent_fragments / progress.total_fragments)
+            )
+            / len(self._latest_version_firmware.files)
+        )
+        self.async_write_ha_state()
+
+    @callback
+    def _reset_progress(self) -> None:
+        """Reset update install progress."""
+        if self._progress_unsub:
+            self._progress_unsub()
+            self._progress_unsub = None
+        self._num_files_installed = 0
+        self._attr_in_progress = False
+        self.async_write_ha_state()
+
     async def _async_update(self, _: HomeAssistant | datetime | None = None) -> None:
         """Update the entity."""
         self._poll_unsub = None
@@ -152,18 +187,29 @@ class ZWaveNodeFirmwareUpdate(UpdateEntity):
         """Install an update."""
         firmware = self._latest_version_firmware
         assert firmware
-        try:
-            for file in firmware.files:
+        self._attr_in_progress = 0
+        self.async_write_ha_state()
+        self._progress_unsub = self.node.on(
+            "firmware update progress", self._update_progress
+        )
+        for file in firmware.files:
+            try:
                 await self.driver.controller.async_begin_ota_firmware_update(
                     self.node, file
                 )
-        except BaseZwaveJSServerError as err:
-            raise HomeAssistantError(err) from err
-        else:
-            self._attr_installed_version = self._attr_latest_version = firmware.version
-            self._latest_version_firmware = None
+            except BaseZwaveJSServerError as err:
+                self._reset_progress()
+                raise HomeAssistantError(err) from err
+            self._num_files_installed += 1
+            self._attr_in_progress = floor(
+                100 * self._num_files_installed / len(firmware.files)
+            )
             self.async_write_ha_state()
 
+        self._attr_installed_version = self._attr_latest_version = firmware.version
+        self._latest_version_firmware = None
+        self._reset_progress()
+
     async def async_poll_value(self, _: bool) -> None:
         """Poll a value."""
         LOGGER.error(
@@ -208,3 +254,7 @@ class ZWaveNodeFirmwareUpdate(UpdateEntity):
         if self._poll_unsub:
             self._poll_unsub()
             self._poll_unsub = None
+
+        if self._progress_unsub:
+            self._progress_unsub()
+            self._progress_unsub = None