From b9132e78b48bc9e0356c3784ae6d357e25cd0db3 Mon Sep 17 00:00:00 2001
From: Franck Nijhof <git@frenck.dev>
Date: Tue, 1 Nov 2022 19:11:50 +0100
Subject: [PATCH] Improve error logging of WebSocket API (#81360)

---
 .../components/websocket_api/connection.py    | 12 ++-
 .../websocket_api/test_connection.py          | 92 +++++++++++++++----
 2 files changed, 82 insertions(+), 22 deletions(-)

diff --git a/homeassistant/components/websocket_api/connection.py b/homeassistant/components/websocket_api/connection.py
index c344e1c6a9f..ab4dda845db 100644
--- a/homeassistant/components/websocket_api/connection.py
+++ b/homeassistant/components/websocket_api/connection.py
@@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any
 import voluptuous as vol
 
 from homeassistant.auth.models import RefreshToken, User
+from homeassistant.components.http import current_request
 from homeassistant.core import Context, HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError, Unauthorized
 
@@ -137,6 +138,13 @@ class ActiveConnection:
             err_message = "Unknown error"
             log_handler = self.logger.exception
 
-        log_handler("Error handling message: %s (%s)", err_message, code)
-
         self.send_message(messages.error_message(msg["id"], code, err_message))
+
+        if code:
+            err_message += f" ({code})"
+        if request := current_request.get():
+            err_message += f" from {request.remote}"
+            if user_agent := request.headers.get("user-agent"):
+                err_message += f" ({user_agent})"
+
+        log_handler("Error handling message: %s", err_message)
diff --git a/tests/components/websocket_api/test_connection.py b/tests/components/websocket_api/test_connection.py
index fd9af99c1a4..8f2cd43fdb8 100644
--- a/tests/components/websocket_api/test_connection.py
+++ b/tests/components/websocket_api/test_connection.py
@@ -1,8 +1,11 @@
 """Test WebSocket Connection class."""
 import asyncio
 import logging
-from unittest.mock import Mock
+from typing import Any
+from unittest.mock import AsyncMock, Mock, patch
 
+from aiohttp.test_utils import make_mocked_request
+import pytest
 import voluptuous as vol
 
 from homeassistant import exceptions
@@ -11,37 +14,86 @@ from homeassistant.components import websocket_api
 from tests.common import MockUser
 
 
-async def test_exception_handling():
-    """Test handling of exceptions."""
-    send_messages = []
-    user = MockUser()
-    refresh_token = Mock()
-    conn = websocket_api.ActiveConnection(
-        logging.getLogger(__name__), None, send_messages.append, user, refresh_token
-    )
-
-    for (exc, code, err) in (
-        (exceptions.Unauthorized(), websocket_api.ERR_UNAUTHORIZED, "Unauthorized"),
+@pytest.mark.parametrize(
+    "exc,code,err,log",
+    [
+        (
+            exceptions.Unauthorized(),
+            websocket_api.ERR_UNAUTHORIZED,
+            "Unauthorized",
+            "Error handling message: Unauthorized (unauthorized) from 127.0.0.42 (Browser)",
+        ),
         (
             vol.Invalid("Invalid something"),
             websocket_api.ERR_INVALID_FORMAT,
             "Invalid something. Got {'id': 5}",
+            "Error handling message: Invalid something. Got {'id': 5} (invalid_format) from 127.0.0.42 (Browser)",
+        ),
+        (
+            asyncio.TimeoutError(),
+            websocket_api.ERR_TIMEOUT,
+            "Timeout",
+            "Error handling message: Timeout (timeout) from 127.0.0.42 (Browser)",
         ),
-        (asyncio.TimeoutError(), websocket_api.ERR_TIMEOUT, "Timeout"),
         (
             exceptions.HomeAssistantError("Failed to do X"),
             websocket_api.ERR_UNKNOWN_ERROR,
             "Failed to do X",
+            "Error handling message: Failed to do X (unknown_error) from 127.0.0.42 (Browser)",
         ),
-        (ValueError("Really bad"), websocket_api.ERR_UNKNOWN_ERROR, "Unknown error"),
         (
-            exceptions.HomeAssistantError(),
+            ValueError("Really bad"),
             websocket_api.ERR_UNKNOWN_ERROR,
             "Unknown error",
+            "Error handling message: Unknown error (unknown_error) from 127.0.0.42 (Browser)",
         ),
-    ):
-        send_messages.clear()
+        (
+            exceptions.HomeAssistantError,
+            websocket_api.ERR_UNKNOWN_ERROR,
+            "Unknown error",
+            "Error handling message: Unknown error (unknown_error) from 127.0.0.42 (Browser)",
+        ),
+    ],
+)
+async def test_exception_handling(
+    caplog: pytest.LogCaptureFixture,
+    exc: Exception,
+    code: str,
+    err: str,
+    log: str,
+):
+    """Test handling of exceptions."""
+    send_messages = []
+    user = MockUser()
+    refresh_token = Mock()
+    current_request = AsyncMock()
+
+    def get_extra_info(key: str) -> Any:
+        if key == "sslcontext":
+            return True
+
+        if key == "peername":
+            return ("127.0.0.42", 8123)
+
+    mocked_transport = Mock()
+    mocked_transport.get_extra_info = get_extra_info
+    mocked_request = make_mocked_request(
+        "GET",
+        "/api/websocket",
+        headers={"Host": "example.com", "User-Agent": "Browser"},
+        transport=mocked_transport,
+    )
+
+    with patch(
+        "homeassistant.components.websocket_api.connection.current_request",
+    ) as current_request:
+        current_request.get.return_value = mocked_request
+        conn = websocket_api.ActiveConnection(
+            logging.getLogger(__name__), None, send_messages.append, user, refresh_token
+        )
+
         conn.async_handle_exception({"id": 5}, exc)
-        assert len(send_messages) == 1
-        assert send_messages[0]["error"]["code"] == code
-        assert send_messages[0]["error"]["message"] == err
+    assert len(send_messages) == 1
+    assert send_messages[0]["error"]["code"] == code
+    assert send_messages[0]["error"]["message"] == err
+    assert log in caplog.text