Lazy load zwave_js platforms when the first entity needs to be created (#49016)

* Lazy load zwave_js platforms when the first entity needs to be created

* switch order to make things easier to understand

* await task instead of using wait_for_done callback

* gather tasks

* switch from asyncio.create_task to hass.async_create_task

* unsubscribe from callbacks before unloading platforms

* Clean up as much as possible during entry unload, even if a platform unload fails

homeassistant/components/zwave_js/__init__.py

@@ -54,11 +54,11 @@ from .const import (
     CONF_USB_PATH,
     CONF_USE_ADDON,
     DATA_CLIENT,
+    DATA_PLATFORM_SETUP,
     DATA_UNSUBSCRIBE,
     DOMAIN,
     EVENT_DEVICE_ADDED_TO_REGISTRY,
     LOGGER,
-    PLATFORMS,
     ZWAVE_JS_NOTIFICATION_EVENT,
     ZWAVE_JS_VALUE_NOTIFICATION_EVENT,
 )
@@ -113,49 +113,69 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     client = ZwaveClient(entry.data[CONF_URL], async_get_clientsession(hass))
     dev_reg = device_registry.async_get(hass)
     ent_reg = entity_registry.async_get(hass)
+    entry_hass_data: dict = hass.data[DOMAIN].setdefault(entry.entry_id, {})
 
-    @callback
+    unsubscribe_callbacks: list[Callable] = []
-    def async_on_node_ready(node: ZwaveNode) -> None:
+    entry_hass_data[DATA_CLIENT] = client
+    entry_hass_data[DATA_UNSUBSCRIBE] = unsubscribe_callbacks
+    entry_hass_data[DATA_PLATFORM_SETUP] = {}
+
+    async def async_on_node_ready(node: ZwaveNode) -> None:
         """Handle node ready event."""
         LOGGER.debug("Processing node %s", node)
 
+        platform_setup_tasks = entry_hass_data[DATA_PLATFORM_SETUP]
+
         # register (or update) node in device registry
         register_node_in_dev_reg(hass, entry, dev_reg, client, node)
 
         # run discovery on all node values and create/update entities
         for disc_info in async_discover_values(node):
-            LOGGER.debug("Discovered entity: %s", disc_info)
-
             # This migration logic was added in 2021.3 to handle a breaking change to
             # the value_id format. Some time in the future, this call (as well as the
             # helper functions) can be removed.
             async_migrate_discovered_value(ent_reg, client, disc_info)
+            if disc_info.platform not in platform_setup_tasks:
+                platform_setup_tasks[disc_info.platform] = hass.async_create_task(
+                    hass.config_entries.async_forward_entry_setup(
+                        entry, disc_info.platform
+                    )
+                )
+
+            await platform_setup_tasks[disc_info.platform]
+
+            LOGGER.debug("Discovered entity: %s", disc_info)
             async_dispatcher_send(
                 hass, f"{DOMAIN}_{entry.entry_id}_add_{disc_info.platform}", disc_info
             )
+
         # add listener for stateless node value notification events
-        node.on(
+        unsubscribe_callbacks.append(
-            "value notification",
+            node.on(
-            lambda event: async_on_value_notification(event["value_notification"]),
+                "value notification",
+                lambda event: async_on_value_notification(event["value_notification"]),
+            )
         )
         # add listener for stateless node notification events
-        node.on(
+        unsubscribe_callbacks.append(
-            "notification", lambda event: async_on_notification(event["notification"])
+            node.on(
+                "notification",
+                lambda event: async_on_notification(event["notification"]),
+            )
         )
 
-    @callback
+    async def async_on_node_added(node: ZwaveNode) -> None:
-    def async_on_node_added(node: ZwaveNode) -> None:
         """Handle node added event."""
         # 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:
-            async_on_node_ready(node)
+            await async_on_node_ready(node)
             return
         # if node is not yet ready, register one-time callback for ready state
         LOGGER.debug("Node added: %s - waiting for it to become ready", node.node_id)
         node.once(
             "ready",
-            lambda event: async_on_node_ready(event["node"]),
+            lambda event: hass.async_create_task(async_on_node_ready(event["node"])),
         )
         # we do submit the node to device registry so user has
         # some visual feedback that something is (in the process of) being added
@@ -234,7 +254,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
         hass.bus.async_fire(ZWAVE_JS_NOTIFICATION_EVENT, event_data)
 
-    entry_hass_data: dict = hass.data[DOMAIN].setdefault(entry.entry_id, {})
     # connect and throw error if connection failed
     try:
         async with timeout(CONNECT_TIMEOUT):
@@ -256,10 +275,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         entry_hass_data[DATA_CONNECT_FAILED_LOGGED] = False
         entry_hass_data[DATA_INVALID_SERVER_VERSION_LOGGED] = False
 
-    unsubscribe_callbacks: list[Callable] = []
-    entry_hass_data[DATA_CLIENT] = client
-    entry_hass_data[DATA_UNSUBSCRIBE] = unsubscribe_callbacks
-
     services = ZWaveServices(hass, ent_reg)
     services.async_register()
 
@@ -268,14 +283,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     async def start_platforms() -> None:
         """Start platforms and perform discovery."""
-        # wait until all required platforms are ready
-        await asyncio.gather(
-            *[
-                hass.config_entries.async_forward_entry_setup(entry, platform)
-                for platform in PLATFORMS
-            ]
-        )
-
         driver_ready = asyncio.Event()
 
         async def handle_ha_shutdown(event: Event) -> None:
@@ -313,17 +320,28 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
                 dev_reg.async_remove_device(device.id)
 
         # run discovery on all ready nodes
-        for node in client.driver.controller.nodes.values():
+        await asyncio.gather(
-            async_on_node_added(node)
+            *[
+                async_on_node_added(node)
+                for node in client.driver.controller.nodes.values()
+            ]
+        )
 
         # listen for new nodes being added to the mesh
-        client.driver.controller.on(
+        unsubscribe_callbacks.append(
-            "node added", lambda event: async_on_node_added(event["node"])
+            client.driver.controller.on(
+                "node added",
+                lambda event: hass.async_create_task(
+                    async_on_node_added(event["node"])
+                ),
+            )
         )
         # listen for nodes being removed from the mesh
         # NOTE: This will not remove nodes that were removed when HA was not running
-        client.driver.controller.on(
+        unsubscribe_callbacks.append(
-            "node removed", lambda event: async_on_node_removed(event["node"])
+            client.driver.controller.on(
+                "node removed", lambda event: async_on_node_removed(event["node"])
+            )
         )
 
     platform_task = hass.async_create_task(start_platforms())
@@ -355,7 +373,7 @@ async def client_listen(
     # All model instances will be replaced when the new state is acquired.
     if should_reload:
         LOGGER.info("Disconnected from server. Reloading integration")
-        asyncio.create_task(hass.config_entries.async_reload(entry.entry_id))
+        hass.async_create_task(hass.config_entries.async_reload(entry.entry_id))
 
 
 async def disconnect_client(
@@ -368,8 +386,13 @@ async def disconnect_client(
     """Disconnect client."""
     listen_task.cancel()
     platform_task.cancel()
+    platform_setup_tasks = (
+        hass.data[DOMAIN].get(entry.entry_id, {}).get(DATA_PLATFORM_SETUP, {}).values()
+    )
+    for task in platform_setup_tasks:
+        task.cancel()
 
-    await asyncio.gather(listen_task, platform_task)
+    await asyncio.gather(listen_task, platform_task, *platform_setup_tasks)
 
     if client.connected:
         await client.disconnect()
@@ -378,22 +401,23 @@ async def disconnect_client(
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
-    unload_ok = all(
-        await asyncio.gather(
-            *[
-                hass.config_entries.async_forward_entry_unload(entry, platform)
-                for platform in PLATFORMS
-            ]
-        )
-    )
-    if not unload_ok:
-        return False
-
     info = hass.data[DOMAIN].pop(entry.entry_id)
 
     for unsub in info[DATA_UNSUBSCRIBE]:
         unsub()
 
+    tasks = []
+    for platform, task in info[DATA_PLATFORM_SETUP].items():
+        if task.done():
+            tasks.append(
+                hass.config_entries.async_forward_entry_unload(entry, platform)
+            )
+        else:
+            task.cancel()
+            tasks.append(task)
+
+    unload_ok = all(await asyncio.gather(*tasks))
+
     if DATA_CLIENT_LISTEN_TASK in info:
         await disconnect_client(
             hass,
@@ -412,7 +436,7 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             LOGGER.error("Failed to stop the Z-Wave JS add-on: %s", err)
             return False
 
-    return True
+    return unload_ok
 
 
 async def async_remove_entry(hass: HomeAssistant, entry: ConfigEntry) -> None:

homeassistant/components/zwave_js/const.py

@@ -8,20 +8,10 @@ CONF_NETWORK_KEY = "network_key"
 CONF_USB_PATH = "usb_path"
 CONF_USE_ADDON = "use_addon"
 DOMAIN = "zwave_js"
-PLATFORMS = [
-    "binary_sensor",
-    "climate",
-    "cover",
-    "fan",
-    "light",
-    "lock",
-    "number",
-    "sensor",
-    "switch",
-]
 
 DATA_CLIENT = "client"
 DATA_UNSUBSCRIBE = "unsubs"
+DATA_PLATFORM_SETUP = "platform_setup"
 
 EVENT_DEVICE_ADDED_TO_REGISTRY = f"{DOMAIN}_device_added_to_registry"