From 99c5087c1e6868a3d7e819e1084728c5dcc558ac Mon Sep 17 00:00:00 2001
From: Raman Gupta <7243222+raman325@users.noreply.github.com>
Date: Wed, 21 Apr 2021 07:37:35 -0400
Subject: [PATCH] Add WS API command to capture zwave_js logs from server
 (#49444)

* Add WS API commands to capture zwave_js logs from server

* register commands

* create a task

* Update homeassistant/components/zwave_js/api.py

Co-authored-by: Paulus Schoutsen <paulus@home-assistant.io>

* Update homeassistant/components/zwave_js/api.py

Co-authored-by: Paulus Schoutsen <paulus@home-assistant.io>

* fix

* fixes and add test

* fix PR on rebase

Co-authored-by: Paulus Schoutsen <paulus@home-assistant.io>
---
 homeassistant/components/zwave_js/api.py | 49 ++++++++++++++++++++++++
 tests/components/zwave_js/test_api.py    | 42 ++++++++++++++++++++
 2 files changed, 91 insertions(+)

diff --git a/homeassistant/components/zwave_js/api.py b/homeassistant/components/zwave_js/api.py
index 9a30d78a07c..be4386e529e 100644
--- a/homeassistant/components/zwave_js/api.py
+++ b/homeassistant/components/zwave_js/api.py
@@ -13,6 +13,7 @@ from zwave_js_server.client import Client
 from zwave_js_server.const import LogLevel
 from zwave_js_server.exceptions import InvalidNewValue, NotFoundError, SetValueFailed
 from zwave_js_server.model.log_config import LogConfig
+from zwave_js_server.model.log_message import LogMessage
 from zwave_js_server.util.node import async_set_config_parameter
 
 from homeassistant.components import websocket_api
@@ -91,6 +92,7 @@ def async_register_api(hass: HomeAssistant) -> None:
     websocket_api.async_register_command(hass, websocket_remove_node)
     websocket_api.async_register_command(hass, websocket_stop_exclusion)
     websocket_api.async_register_command(hass, websocket_refresh_node_info)
+    websocket_api.async_register_command(hass, websocket_subscribe_logs)
     websocket_api.async_register_command(hass, websocket_update_log_config)
     websocket_api.async_register_command(hass, websocket_get_log_config)
     websocket_api.async_register_command(hass, websocket_get_config_parameters)
@@ -517,6 +519,53 @@ def filename_is_present_if_logging_to_file(obj: dict) -> dict:
     return obj
 
 
+@websocket_api.require_admin  # type: ignore
+@websocket_api.async_response
+@websocket_api.websocket_command(
+    {
+        vol.Required(TYPE): "zwave_js/subscribe_logs",
+        vol.Required(ENTRY_ID): str,
+    }
+)
+@async_get_entry
+async def websocket_subscribe_logs(
+    hass: HomeAssistant,
+    connection: ActiveConnection,
+    msg: dict,
+    entry: ConfigEntry,
+    client: Client,
+) -> None:
+    """Subscribe to log message events from the server."""
+    driver = client.driver
+
+    @callback
+    def async_cleanup() -> None:
+        """Remove signal listeners."""
+        hass.async_create_task(driver.async_stop_listening_logs())
+        unsub()
+
+    @callback
+    def forward_event(event: dict) -> None:
+        log_msg: LogMessage = event["log_message"]
+        connection.send_message(
+            websocket_api.event_message(
+                msg[ID],
+                {
+                    "timestamp": log_msg.timestamp,
+                    "level": log_msg.level,
+                    "primary_tags": log_msg.primary_tags,
+                    "message": log_msg.formatted_message,
+                },
+            )
+        )
+
+    unsub = driver.on("logging", forward_event)
+    connection.subscriptions[msg["id"]] = async_cleanup
+
+    await driver.async_start_listening_logs()
+    connection.send_result(msg[ID])
+
+
 @websocket_api.require_admin  # type: ignore
 @websocket_api.async_response
 @websocket_api.websocket_command(
diff --git a/tests/components/zwave_js/test_api.py b/tests/components/zwave_js/test_api.py
index 525f97da681..3fb57a366aa 100644
--- a/tests/components/zwave_js/test_api.py
+++ b/tests/components/zwave_js/test_api.py
@@ -445,6 +445,48 @@ async def test_dump_view_invalid_entry_id(integration, hass_client):
     assert resp.status == 400
 
 
+async def test_subscribe_logs(hass, integration, client, hass_ws_client):
+    """Test the subscribe_logs websocket command."""
+    entry = integration
+    ws_client = await hass_ws_client(hass)
+
+    client.async_send_command.return_value = {}
+
+    await ws_client.send_json(
+        {ID: 1, TYPE: "zwave_js/subscribe_logs", ENTRY_ID: entry.entry_id}
+    )
+
+    msg = await ws_client.receive_json()
+    assert msg["success"]
+
+    event = Event(
+        type="logging",
+        data={
+            "source": "driver",
+            "event": "logging",
+            "message": "test",
+            "formattedMessage": "test",
+            "direction": ">",
+            "level": "debug",
+            "primaryTags": "tag",
+            "secondaryTags": "tag2",
+            "secondaryTagPadding": 0,
+            "multiline": False,
+            "timestamp": "time",
+            "label": "label",
+        },
+    )
+    client.driver.receive_event(event)
+
+    msg = await ws_client.receive_json()
+    assert msg["event"] == {
+        "message": ["test"],
+        "level": "debug",
+        "primary_tags": "tag",
+        "timestamp": "time",
+    }
+
+
 async def test_update_log_config(hass, client, integration, hass_ws_client):
     """Test that the update_log_config WS API call works and that schema validation works."""
     entry = integration