From 394442e8a906a327eea1d3713f6d10eda043fdc7 Mon Sep 17 00:00:00 2001
From: Raman Gupta <7243222+raman325@users.noreply.github.com>
Date: Wed, 1 Jun 2022 00:42:07 -0400
Subject: [PATCH] Use device_id for zwave_js/replace_failed_node command
 (#72785)

---
 homeassistant/components/zwave_js/api.py | 15 +++++-------
 tests/components/zwave_js/test_api.py    | 31 ++++++++++--------------
 2 files changed, 19 insertions(+), 27 deletions(-)

diff --git a/homeassistant/components/zwave_js/api.py b/homeassistant/components/zwave_js/api.py
index f9f1ad40f9a..e6ce08a78a5 100644
--- a/homeassistant/components/zwave_js/api.py
+++ b/homeassistant/components/zwave_js/api.py
@@ -1083,8 +1083,7 @@ async def websocket_remove_node(
 @websocket_api.websocket_command(
     {
         vol.Required(TYPE): "zwave_js/replace_failed_node",
-        vol.Required(ENTRY_ID): str,
-        vol.Required(NODE_ID): int,
+        vol.Required(DEVICE_ID): str,
         vol.Optional(INCLUSION_STRATEGY, default=InclusionStrategy.DEFAULT): vol.All(
             vol.Coerce(int),
             vol.In(
@@ -1107,18 +1106,16 @@ async def websocket_remove_node(
 )
 @websocket_api.async_response
 @async_handle_failed_command
-@async_get_entry
+@async_get_node
 async def websocket_replace_failed_node(
     hass: HomeAssistant,
     connection: ActiveConnection,
     msg: dict,
-    entry: ConfigEntry,
-    client: Client,
-    driver: Driver,
+    node: Node,
 ) -> None:
     """Replace a failed node with a new node."""
-    controller = driver.controller
-    node_id = msg[NODE_ID]
+    assert node.client.driver
+    controller = node.client.driver.controller
     inclusion_strategy = InclusionStrategy(msg[INCLUSION_STRATEGY])
     force_security = msg.get(FORCE_SECURITY)
     provisioning = (
@@ -1232,7 +1229,7 @@ async def websocket_replace_failed_node(
 
     try:
         result = await controller.async_replace_failed_node(
-            controller.nodes[node_id],
+            node,
             INCLUSION_STRATEGY_NOT_SMART_START[inclusion_strategy.value],
             force_security=force_security,
             provisioning=provisioning,
diff --git a/tests/components/zwave_js/test_api.py b/tests/components/zwave_js/test_api.py
index 7f64ed6d87d..da8cfad9624 100644
--- a/tests/components/zwave_js/test_api.py
+++ b/tests/components/zwave_js/test_api.py
@@ -1511,7 +1511,7 @@ async def test_replace_failed_node(
     dev_reg = dr.async_get(hass)
 
     # Create device registry entry for mock node
-    dev_reg.async_get_or_create(
+    device = dev_reg.async_get_or_create(
         config_entry_id=entry.entry_id,
         identifiers={(DOMAIN, "3245146787-67")},
         name="Node 67",
@@ -1526,8 +1526,7 @@ async def test_replace_failed_node(
         {
             ID: 1,
             TYPE: "zwave_js/replace_failed_node",
-            ENTRY_ID: entry.entry_id,
-            NODE_ID: 67,
+            DEVICE_ID: device.id,
             INCLUSION_STRATEGY: InclusionStrategy.DEFAULT.value,
         }
     )
@@ -1607,10 +1606,12 @@ async def test_replace_failed_node(
     assert msg["event"]["event"] == "node removed"
 
     # Verify device was removed from device registry
-    device = dev_reg.async_get_device(
-        identifiers={(DOMAIN, "3245146787-67")},
+    assert (
+        dev_reg.async_get_device(
+            identifiers={(DOMAIN, "3245146787-67")},
+        )
+        is None
     )
-    assert device is None
 
     client.driver.receive_event(nortek_thermostat_added_event)
     msg = await ws_client.receive_json()
@@ -1686,8 +1687,7 @@ async def test_replace_failed_node(
         {
             ID: 2,
             TYPE: "zwave_js/replace_failed_node",
-            ENTRY_ID: entry.entry_id,
-            NODE_ID: 67,
+            DEVICE_ID: device.id,
             INCLUSION_STRATEGY: InclusionStrategy.SECURITY_S2.value,
             PLANNED_PROVISIONING_ENTRY: {
                 DSK: "test",
@@ -1719,8 +1719,7 @@ async def test_replace_failed_node(
         {
             ID: 3,
             TYPE: "zwave_js/replace_failed_node",
-            ENTRY_ID: entry.entry_id,
-            NODE_ID: 67,
+            DEVICE_ID: device.id,
             INCLUSION_STRATEGY: InclusionStrategy.SECURITY_S2.value,
             QR_PROVISIONING_INFORMATION: {
                 VERSION: 0,
@@ -1772,8 +1771,7 @@ async def test_replace_failed_node(
         {
             ID: 4,
             TYPE: "zwave_js/replace_failed_node",
-            ENTRY_ID: entry.entry_id,
-            NODE_ID: 67,
+            DEVICE_ID: device.id,
             INCLUSION_STRATEGY: InclusionStrategy.SECURITY_S2.value,
             QR_CODE_STRING: "90testtesttesttesttesttesttesttesttesttesttesttesttest",
         }
@@ -1800,8 +1798,7 @@ async def test_replace_failed_node(
         {
             ID: 6,
             TYPE: "zwave_js/replace_failed_node",
-            ENTRY_ID: entry.entry_id,
-            NODE_ID: 67,
+            DEVICE_ID: device.id,
             INCLUSION_STRATEGY: InclusionStrategy.DEFAULT.value,
             QR_CODE_STRING: "90testtesttesttesttesttesttesttesttesttesttesttesttest",
         }
@@ -1821,8 +1818,7 @@ async def test_replace_failed_node(
             {
                 ID: 7,
                 TYPE: "zwave_js/replace_failed_node",
-                ENTRY_ID: entry.entry_id,
-                NODE_ID: 67,
+                DEVICE_ID: device.id,
             }
         )
         msg = await ws_client.receive_json()
@@ -1839,8 +1835,7 @@ async def test_replace_failed_node(
         {
             ID: 8,
             TYPE: "zwave_js/replace_failed_node",
-            ENTRY_ID: entry.entry_id,
-            NODE_ID: 67,
+            DEVICE_ID: device.id,
         }
     )
     msg = await ws_client.receive_json()