From 4ff1f5c7886b7ac2e4a6b582b4a2df6f380e2e23 Mon Sep 17 00:00:00 2001
From: Raman Gupta <7243222+raman325@users.noreply.github.com>
Date: Sat, 19 Feb 2022 05:16:14 -0500
Subject: [PATCH] Create zwave_js ping button at the right time (#66848)

* Create zwave_js ping button at the right time

* fix tests
---
 homeassistant/components/zwave_js/__init__.py | 15 +++++++--------
 tests/components/zwave_js/test_init.py        |  8 ++++----
 2 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/homeassistant/components/zwave_js/__init__.py b/homeassistant/components/zwave_js/__init__.py
index 0e1c6445a9f..1682edb66a3 100644
--- a/homeassistant/components/zwave_js/__init__.py
+++ b/homeassistant/components/zwave_js/__init__.py
@@ -261,12 +261,6 @@ async def async_setup_entry(  # noqa: C901
             )
         )
 
-        # Create a ping button for each device
-        await async_setup_platform(BUTTON_DOMAIN)
-        async_dispatcher_send(
-            hass, f"{DOMAIN}_{entry.entry_id}_add_ping_button_entity", node
-        )
-
         # add listeners to handle new values that get added later
         for event in ("value added", "value updated", "metadata updated"):
             entry.async_on_unload(
@@ -295,13 +289,18 @@ async def async_setup_entry(  # noqa: C901
 
     async def async_on_node_added(node: ZwaveNode) -> None:
         """Handle node added event."""
-        await async_setup_platform(SENSOR_DOMAIN)
-
         # Create a node status sensor for each device
+        await async_setup_platform(SENSOR_DOMAIN)
         async_dispatcher_send(
             hass, f"{DOMAIN}_{entry.entry_id}_add_node_status_sensor", node
         )
 
+        # Create a ping button for each device
+        await async_setup_platform(BUTTON_DOMAIN)
+        async_dispatcher_send(
+            hass, f"{DOMAIN}_{entry.entry_id}_add_ping_button_entity", node
+        )
+
         # we only want to run discovery when the node has reached ready state,
         # otherwise we'll have all kinds of missing info issues.
         if node.ready:
diff --git a/tests/components/zwave_js/test_init.py b/tests/components/zwave_js/test_init.py
index 3780b2654f9..1003316f1e5 100644
--- a/tests/components/zwave_js/test_init.py
+++ b/tests/components/zwave_js/test_init.py
@@ -211,8 +211,8 @@ async def test_on_node_added_not_ready(
     client.driver.receive_event(event)
     await hass.async_block_till_done()
 
-    # the only entity is the node status sensor
-    assert len(hass.states.async_all()) == 1
+    # the only entities are the node status sensor and ping button
+    assert len(hass.states.async_all()) == 2
 
     device = dev_reg.async_get_device(identifiers={(DOMAIN, device_id)})
     assert device
@@ -254,8 +254,8 @@ async def test_existing_node_not_ready(hass, zp3111_not_ready, client, integrati
     assert not device.model
     assert not device.sw_version
 
-    # the only entity is the node status sensor
-    assert len(hass.states.async_all()) == 1
+    # the only entities are the node status sensor and ping button
+    assert len(hass.states.async_all()) == 2
 
     device = dev_reg.async_get_device(identifiers={(DOMAIN, device_id)})
     assert device