From cad281b3265c2a023f3a5c1f49805bc2723c6c80 Mon Sep 17 00:00:00 2001
From: Charles Garwood <cgarwood@gmail.com>
Date: Wed, 21 Apr 2021 07:35:16 -0400
Subject: [PATCH] Add subscription for Z-Wave JS node re-interview status
 (#49024)

* Add subscription for interview status

* update test

* forward stage completed event

* add additional test

* additional tests

* return earlier
---
 homeassistant/components/zwave_js/api.py | 35 ++++++++++++++++++--
 tests/components/zwave_js/test_api.py    | 42 ++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 3 deletions(-)

diff --git a/homeassistant/components/zwave_js/api.py b/homeassistant/components/zwave_js/api.py
index fc4f16bda33..9a30d78a07c 100644
--- a/homeassistant/components/zwave_js/api.py
+++ b/homeassistant/components/zwave_js/api.py
@@ -367,14 +367,43 @@ async def websocket_refresh_node_info(
 ) -> None:
     """Re-interview a node."""
     node_id = msg[NODE_ID]
-    node = client.driver.controller.nodes.get(node_id)
+    controller = client.driver.controller
+    node = controller.nodes.get(node_id)
 
     if node is None:
         connection.send_error(msg[ID], ERR_NOT_FOUND, f"Node {node_id} not found")
         return
 
-    await node.async_refresh_info()
-    connection.send_result(msg[ID])
+    @callback
+    def async_cleanup() -> None:
+        """Remove signal listeners."""
+        for unsub in unsubs:
+            unsub()
+
+    @callback
+    def forward_event(event: dict) -> None:
+        connection.send_message(
+            websocket_api.event_message(msg[ID], {"event": event["event"]})
+        )
+
+    @callback
+    def forward_stage(event: dict) -> None:
+        connection.send_message(
+            websocket_api.event_message(
+                msg[ID], {"event": event["event"], "stage": event["stageName"]}
+            )
+        )
+
+    connection.subscriptions[msg["id"]] = async_cleanup
+    unsubs = [
+        node.on("interview started", forward_event),
+        node.on("interview completed", forward_event),
+        node.on("interview stage completed", forward_stage),
+        node.on("interview failed", forward_event),
+    ]
+
+    result = await node.async_refresh_info()
+    connection.send_result(msg[ID], result)
 
 
 @websocket_api.require_admin  # type:ignore
diff --git a/tests/components/zwave_js/test_api.py b/tests/components/zwave_js/test_api.py
index edd711b07d1..525f97da681 100644
--- a/tests/components/zwave_js/test_api.py
+++ b/tests/components/zwave_js/test_api.py
@@ -250,6 +250,48 @@ async def test_refresh_node_info(
     assert args["command"] == "node.refresh_info"
     assert args["nodeId"] == 52
 
+    event = Event(
+        type="interview started",
+        data={"source": "node", "event": "interview started", "nodeId": 52},
+    )
+    client.driver.receive_event(event)
+
+    msg = await ws_client.receive_json()
+    assert msg["event"]["event"] == "interview started"
+
+    event = Event(
+        type="interview stage completed",
+        data={
+            "source": "node",
+            "event": "interview stage completed",
+            "stageName": "NodeInfo",
+            "nodeId": 52,
+        },
+    )
+    client.driver.receive_event(event)
+
+    msg = await ws_client.receive_json()
+    assert msg["event"]["event"] == "interview stage completed"
+    assert msg["event"]["stage"] == "NodeInfo"
+
+    event = Event(
+        type="interview completed",
+        data={"source": "node", "event": "interview completed", "nodeId": 52},
+    )
+    client.driver.receive_event(event)
+
+    msg = await ws_client.receive_json()
+    assert msg["event"]["event"] == "interview completed"
+
+    event = Event(
+        type="interview failed",
+        data={"source": "node", "event": "interview failed", "nodeId": 52},
+    )
+    client.driver.receive_event(event)
+
+    msg = await ws_client.receive_json()
+    assert msg["event"]["event"] == "interview failed"
+
     client.async_send_command_no_wait.reset_mock()
 
     await ws_client.send_json(