From 30f84f55a4c3f45e0df3fe239461169ac499d1ac Mon Sep 17 00:00:00 2001
From: Petar Petrov <MindFreeze@users.noreply.github.com>
Date: Fri, 6 Dec 2024 10:35:48 +0200
Subject: [PATCH] Handle Z-Wave JS S2 inclusion via Inclusion Controller
 (#132073)

* ZwaveJS: Handle S2 inclusion via Inclusion Controller

* improved tests
---
 homeassistant/components/zwave_js/api.py | 35 +++++++++++++
 tests/components/zwave_js/test_api.py    | 62 ++++++++++++++++++++++++
 2 files changed, 97 insertions(+)

diff --git a/homeassistant/components/zwave_js/api.py b/homeassistant/components/zwave_js/api.py
index 88f8f25c8e2..1a1cd6ae9c1 100644
--- a/homeassistant/components/zwave_js/api.py
+++ b/homeassistant/components/zwave_js/api.py
@@ -396,6 +396,7 @@ def async_register_api(hass: HomeAssistant) -> None:
     websocket_api.async_register_command(hass, websocket_node_alerts)
     websocket_api.async_register_command(hass, websocket_add_node)
     websocket_api.async_register_command(hass, websocket_cancel_secure_bootstrap_s2)
+    websocket_api.async_register_command(hass, websocket_subscribe_s2_inclusion)
     websocket_api.async_register_command(hass, websocket_grant_security_classes)
     websocket_api.async_register_command(hass, websocket_validate_dsk_and_enter_pin)
     websocket_api.async_register_command(hass, websocket_provision_smart_start_node)
@@ -863,6 +864,40 @@ async def websocket_cancel_secure_bootstrap_s2(
     connection.send_result(msg[ID])
 
 
+@websocket_api.require_admin
+@websocket_api.websocket_command(
+    {
+        vol.Required(TYPE): "zwave_js/subscribe_s2_inclusion",
+        vol.Required(ENTRY_ID): str,
+    }
+)
+@websocket_api.async_response
+@async_handle_failed_command
+@async_get_entry
+async def websocket_subscribe_s2_inclusion(
+    hass: HomeAssistant,
+    connection: ActiveConnection,
+    msg: dict[str, Any],
+    entry: ConfigEntry,
+    client: Client,
+    driver: Driver,
+) -> None:
+    """Subscribe to S2 inclusion initiated by the controller."""
+
+    @callback
+    def forward_dsk(event: dict) -> None:
+        connection.send_message(
+            websocket_api.event_message(
+                msg[ID], {"event": event["event"], "dsk": event["dsk"]}
+            )
+        )
+
+    unsub = driver.controller.on("validate dsk and enter pin", forward_dsk)
+    connection.subscriptions[msg["id"]] = unsub
+    msg[DATA_UNSUBSCRIBE] = [unsub]
+    connection.send_result(msg[ID])
+
+
 @websocket_api.require_admin
 @websocket_api.websocket_command(
     {
diff --git a/tests/components/zwave_js/test_api.py b/tests/components/zwave_js/test_api.py
index 3761ba6eaa6..a3f70e92dcf 100644
--- a/tests/components/zwave_js/test_api.py
+++ b/tests/components/zwave_js/test_api.py
@@ -5261,3 +5261,65 @@ async def test_cancel_secure_bootstrap_s2(
 
     assert not msg["success"]
     assert msg["error"]["code"] == ERR_NOT_FOUND
+
+
+async def test_subscribe_s2_inclusion(
+    hass: HomeAssistant, integration, client, hass_ws_client: WebSocketGenerator
+) -> None:
+    """Test the subscribe_s2_inclusion websocket command."""
+    entry = integration
+    ws_client = await hass_ws_client(hass)
+
+    await ws_client.send_json_auto_id(
+        {
+            TYPE: "zwave_js/subscribe_s2_inclusion",
+            ENTRY_ID: entry.entry_id,
+        }
+    )
+
+    msg = await ws_client.receive_json()
+    assert msg["success"]
+    assert msg["result"] is None
+
+    # Test receiving DSK request event
+    event = Event(
+        type="validate dsk and enter pin",
+        data={
+            "source": "controller",
+            "event": "validate dsk and enter pin",
+            "dsk": "test_dsk",
+        },
+    )
+    client.driver.receive_event(event)
+
+    msg = await ws_client.receive_json()
+    assert msg["event"] == {
+        "event": "validate dsk and enter pin",
+        "dsk": "test_dsk",
+    }
+
+    # Test sending command with not loaded entry fails
+    await hass.config_entries.async_unload(entry.entry_id)
+    await hass.async_block_till_done()
+
+    await ws_client.send_json_auto_id(
+        {
+            TYPE: "zwave_js/subscribe_s2_inclusion",
+            ENTRY_ID: entry.entry_id,
+        }
+    )
+    msg = await ws_client.receive_json()
+
+    assert not msg["success"]
+    assert msg["error"]["code"] == ERR_NOT_LOADED
+
+    # Test invalid config entry id
+    await ws_client.send_json_auto_id(
+        {
+            TYPE: "zwave_js/subscribe_s2_inclusion",
+            ENTRY_ID: "INVALID",
+        }
+    )
+    msg = await ws_client.receive_json()
+    assert not msg["success"]
+    assert msg["error"]["code"] == ERR_NOT_FOUND