From 0d40ba463ead067f300b9f68ec278e232e125767 Mon Sep 17 00:00:00 2001
From: Raman Gupta <7243222+raman325@users.noreply.github.com>
Date: Mon, 14 Jun 2021 13:31:44 -0400
Subject: [PATCH] Create zwave_js node status sensor when the node is added
 (#51850)

* Create node status sensor when the node is added

* Handle race condition

* reduce repeat code
---
 homeassistant/components/zwave_js/__init__.py | 32 +++++++++++--------
 1 file changed, 19 insertions(+), 13 deletions(-)

diff --git a/homeassistant/components/zwave_js/__init__.py b/homeassistant/components/zwave_js/__init__.py
index 2878580e30c..fee4da743c8 100644
--- a/homeassistant/components/zwave_js/__init__.py
+++ b/homeassistant/components/zwave_js/__init__.py
@@ -179,19 +179,6 @@ async def async_setup_entry(  # noqa: C901
             if disc_info.assumed_state:
                 value_updates_disc_info.append(disc_info)
 
-        # We need to set up the sensor platform if it hasn't already been setup in
-        # order to create the node status sensor
-        if SENSOR_DOMAIN not in platform_setup_tasks:
-            platform_setup_tasks[SENSOR_DOMAIN] = hass.async_create_task(
-                hass.config_entries.async_forward_entry_setup(entry, SENSOR_DOMAIN)
-            )
-            await platform_setup_tasks[SENSOR_DOMAIN]
-
-        # Create a node status sensor for each device
-        async_dispatcher_send(
-            hass, f"{DOMAIN}_{entry.entry_id}_add_node_status_sensor", node
-        )
-
         # add listener for value updated events if necessary
         if value_updates_disc_info:
             unsubscribe_callbacks.append(
@@ -220,6 +207,25 @@ async def async_setup_entry(  # noqa: C901
 
     async def async_on_node_added(node: ZwaveNode) -> None:
         """Handle node added event."""
+        platform_setup_tasks = entry_hass_data[DATA_PLATFORM_SETUP]
+
+        # We need to set up the sensor platform if it hasn't already been setup in
+        # order to create the node status sensor
+        if SENSOR_DOMAIN not in platform_setup_tasks:
+            platform_setup_tasks[SENSOR_DOMAIN] = hass.async_create_task(
+                hass.config_entries.async_forward_entry_setup(entry, SENSOR_DOMAIN)
+            )
+
+        # This guard ensures that concurrent runs of this function all await the
+        # platform setup task
+        if not platform_setup_tasks[SENSOR_DOMAIN].done():
+            await platform_setup_tasks[SENSOR_DOMAIN]
+
+        # Create a node status sensor for each device
+        async_dispatcher_send(
+            hass, f"{DOMAIN}_{entry.entry_id}_add_node_status_sensor", 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: