From 9ee7e55f10961ac8703e1d4c17106cbbe79b1f19 Mon Sep 17 00:00:00 2001
From: Raman Gupta <7243222+raman325@users.noreply.github.com>
Date: Fri, 23 Jul 2021 10:34:49 -0400
Subject: [PATCH] Send initial status in zwave_js WS API cmds to subscribe to
 updates (#53386)

---
 homeassistant/components/zwave_js/api.py | 12 +++-
 tests/components/zwave_js/test_api.py    | 73 ++++++++++++++++++++++++
 2 files changed, 83 insertions(+), 2 deletions(-)

diff --git a/homeassistant/components/zwave_js/api.py b/homeassistant/components/zwave_js/api.py
index 1e1dc2382fc..379376bb98d 100644
--- a/homeassistant/components/zwave_js/api.py
+++ b/homeassistant/components/zwave_js/api.py
@@ -796,7 +796,7 @@ async def websocket_subscribe_heal_network_progress(
         controller.on("heal network done", partial(forward_event, "result")),
     ]
 
-    connection.send_result(msg[ID])
+    connection.send_result(msg[ID], controller.heal_network_progress)
 
 
 @websocket_api.require_admin
@@ -1390,7 +1390,15 @@ async def websocket_subscribe_firmware_update_status(
     ]
     connection.subscriptions[msg["id"]] = async_cleanup
 
-    connection.send_result(msg[ID])
+    result = (
+        {
+            "sent_fragments": node.firmware_update_progress.sent_fragments,
+            "total_fragments": node.firmware_update_progress.total_fragments,
+        }
+        if node.firmware_update_progress
+        else None
+    )
+    connection.send_result(msg[ID], result)
 
 
 class FirmwareUploadView(HomeAssistantView):
diff --git a/tests/components/zwave_js/test_api.py b/tests/components/zwave_js/test_api.py
index 80fe6da90f5..a96e76be865 100644
--- a/tests/components/zwave_js/test_api.py
+++ b/tests/components/zwave_js/test_api.py
@@ -929,6 +929,7 @@ async def test_subscribe_heal_network_progress(
 
     msg = await ws_client.receive_json()
     assert msg["success"]
+    assert msg["result"] is None
 
     # Fire heal network progress
     event = Event(
@@ -961,6 +962,39 @@ async def test_subscribe_heal_network_progress(
     assert msg["error"]["code"] == ERR_NOT_LOADED
 
 
+async def test_subscribe_heal_network_progress_initial_value(
+    hass, integration, client, hass_ws_client
+):
+    """Test subscribe_heal_network_progress command when heal network in progress."""
+    entry = integration
+    ws_client = await hass_ws_client(hass)
+
+    assert not client.driver.controller.heal_network_progress
+
+    # Fire heal network progress before sending heal network progress command
+    event = Event(
+        "heal network progress",
+        {
+            "source": "controller",
+            "event": "heal network progress",
+            "progress": {67: "pending"},
+        },
+    )
+    client.driver.controller.receive_event(event)
+
+    await ws_client.send_json(
+        {
+            ID: 3,
+            TYPE: "zwave_js/subscribe_heal_network_progress",
+            ENTRY_ID: entry.entry_id,
+        }
+    )
+
+    msg = await ws_client.receive_json()
+    assert msg["success"]
+    assert msg["result"] == {"67": "pending"}
+
+
 async def test_stop_healing_network(
     hass,
     integration,
@@ -2403,6 +2437,7 @@ async def test_subscribe_firmware_update_status(
 
     msg = await ws_client.receive_json()
     assert msg["success"]
+    assert msg["result"] is None
 
     event = Event(
         type="firmware update progress",
@@ -2443,6 +2478,44 @@ async def test_subscribe_firmware_update_status(
     }
 
 
+async def test_subscribe_firmware_update_status_initial_value(
+    hass, integration, multisensor_6, client, hass_ws_client
+):
+    """Test subscribe_firmware_update_status websocket command with in progress update."""
+    entry = integration
+    ws_client = await hass_ws_client(hass)
+
+    assert multisensor_6.firmware_update_progress is None
+
+    # Send a firmware update progress event before the WS command
+    event = Event(
+        type="firmware update progress",
+        data={
+            "source": "node",
+            "event": "firmware update progress",
+            "nodeId": multisensor_6.node_id,
+            "sentFragments": 1,
+            "totalFragments": 10,
+        },
+    )
+    multisensor_6.receive_event(event)
+
+    client.async_send_command_no_wait.return_value = {}
+
+    await ws_client.send_json(
+        {
+            ID: 1,
+            TYPE: "zwave_js/subscribe_firmware_update_status",
+            ENTRY_ID: entry.entry_id,
+            NODE_ID: multisensor_6.node_id,
+        }
+    )
+
+    msg = await ws_client.receive_json()
+    assert msg["success"]
+    assert msg["result"] == {"sent_fragments": 1, "total_fragments": 10}
+
+
 async def test_subscribe_firmware_update_status_failures(
     hass, integration, multisensor_6, client, hass_ws_client
 ):