From ec3ee7f02cfa8d9a2b0424f2c865c1196429698d Mon Sep 17 00:00:00 2001
From: Raman Gupta <7243222+raman325@users.noreply.github.com>
Date: Tue, 24 Oct 2023 21:31:03 -0400
Subject: [PATCH] Update zwave_js/hard_reset_controller WS cmd (#102280)

---
 homeassistant/components/zwave_js/api.py | 21 ++++++++++++++++++++-
 tests/components/zwave_js/conftest.py    | 14 +++++++++++---
 tests/components/zwave_js/test_api.py    | 18 ++++++++++++++++--
 3 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/homeassistant/components/zwave_js/api.py b/homeassistant/components/zwave_js/api.py
index 0e7c36c479d..a917aa44889 100644
--- a/homeassistant/components/zwave_js/api.py
+++ b/homeassistant/components/zwave_js/api.py
@@ -2463,5 +2463,24 @@ async def websocket_hard_reset_controller(
     driver: Driver,
 ) -> None:
     """Hard reset controller."""
+
+    @callback
+    def async_cleanup() -> None:
+        """Remove signal listeners."""
+        for unsub in unsubs:
+            unsub()
+        unsubs.clear()
+
+    @callback
+    def _handle_device_added(device: dr.DeviceEntry) -> None:
+        """Handle device is added."""
+        if entry.entry_id in device.config_entries:
+            connection.send_result(msg[ID], device.id)
+            async_cleanup()
+
+    msg[DATA_UNSUBSCRIBE] = unsubs = [
+        async_dispatcher_connect(
+            hass, EVENT_DEVICE_ADDED_TO_REGISTRY, _handle_device_added
+        )
+    ]
     await driver.async_hard_reset()
-    connection.send_result(msg[ID])
diff --git a/tests/components/zwave_js/conftest.py b/tests/components/zwave_js/conftest.py
index db5495bce01..534f2fd2457 100644
--- a/tests/components/zwave_js/conftest.py
+++ b/tests/components/zwave_js/conftest.py
@@ -677,9 +677,19 @@ def central_scene_node_state_fixture():
 # model fixtures
 
 
+@pytest.fixture(name="listen_block")
+def mock_listen_block_fixture():
+    """Mock a listen block."""
+    return asyncio.Event()
+
+
 @pytest.fixture(name="client")
 def mock_client_fixture(
-    controller_state, controller_node_state, version_state, log_config_state
+    controller_state,
+    controller_node_state,
+    version_state,
+    log_config_state,
+    listen_block,
 ):
     """Mock a client."""
     with patch(
@@ -693,9 +703,7 @@ def mock_client_fixture(
 
         async def listen(driver_ready: asyncio.Event) -> None:
             driver_ready.set()
-            listen_block = asyncio.Event()
             await listen_block.wait()
-            pytest.fail("Listen wasn't canceled!")
 
         async def disconnect():
             client.connected = False
diff --git a/tests/components/zwave_js/test_api.py b/tests/components/zwave_js/test_api.py
index 4ff7c481e37..9c4a6339a78 100644
--- a/tests/components/zwave_js/test_api.py
+++ b/tests/components/zwave_js/test_api.py
@@ -4653,12 +4653,21 @@ async def test_subscribe_node_statistics(
 
 
 async def test_hard_reset_controller(
-    hass: HomeAssistant, client, integration, hass_ws_client: WebSocketGenerator
+    hass: HomeAssistant,
+    client,
+    integration,
+    listen_block,
+    hass_ws_client: WebSocketGenerator,
 ) -> None:
     """Test that the hard_reset_controller WS API call works."""
     entry = integration
     ws_client = await hass_ws_client(hass)
 
+    dev_reg = dr.async_get(hass)
+    device = dev_reg.async_get_device(
+        identifiers={get_device_id(client.driver, client.driver.controller.nodes[1])}
+    )
+
     client.async_send_command.return_value = {}
     await ws_client.send_json(
         {
@@ -4667,8 +4676,13 @@ async def test_hard_reset_controller(
             ENTRY_ID: entry.entry_id,
         }
     )
+
+    listen_block.set()
+    listen_block.clear()
+    await hass.async_block_till_done()
+
     msg = await ws_client.receive_json()
-    assert msg["result"] is None
+    assert msg["result"] == device.id
     assert msg["success"]
 
     assert len(client.async_send_command.call_args_list) == 1