Handle WebSocket client disconnect during prepare (#124173)

homeassistant/components/websocket_api/http.py

@@ -306,6 +306,13 @@ class WebSocketHandler:
         try:
             async with asyncio.timeout(10):
                 await wsock.prepare(request)
+        except ConnectionResetError:
+            # Likely the client disconnected before we prepared the websocket
+            logger.debug(
+                "%s: Connection reset by peer while preparing WebSocket",
+                self.description,
+            )
+            return wsock
         except TimeoutError:
             logger.warning("Timeout preparing request from %s", request.remote)
             return wsock

tests/components/websocket_api/test_http.py

@@ -363,12 +363,12 @@ async def test_non_json_message(
     assert "bad=<object" in caplog.text
 
 
-async def test_prepare_fail(
+async def test_prepare_fail_timeout(
     hass: HomeAssistant,
     hass_ws_client: WebSocketGenerator,
     caplog: pytest.LogCaptureFixture,
 ) -> None:
-    """Test failing to prepare."""
+    """Test failing to prepare due to timeout."""
     with (
         patch(
             "homeassistant.components.websocket_api.http.web.WebSocketResponse.prepare",
@@ -381,6 +381,24 @@ async def test_prepare_fail(
     assert "Timeout preparing request" in caplog.text
 
 
+async def test_prepare_fail_connection_reset(
+    hass: HomeAssistant,
+    hass_ws_client: WebSocketGenerator,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test failing to prepare due to connection reset."""
+    with (
+        patch(
+            "homeassistant.components.websocket_api.http.web.WebSocketResponse.prepare",
+            side_effect=(ConnectionResetError, web.WebSocketResponse.prepare),
+        ),
+        pytest.raises(WSServerHandshakeError),
+    ):
+        await hass_ws_client(hass)
+
+    assert "Connection reset by peer while preparing WebSocket" in caplog.text
+
+
 async def test_enable_coalesce(
     hass: HomeAssistant,
     hass_ws_client: WebSocketGenerator,