evlli/hass-core
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Handle prepare timeout in websocket API (#55989)

Browse source
This commit is contained in:
Paulus Schoutsen 2021-10-08 22:12:06 -07:00 • committed by GitHub
parent d55a7e5cc7
commit 6d0da631bf
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 22 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
homeassistant/components/websocket_api/http.py
View file

@ -61,7 +61,7 @@ class WebSocketHandler:
"""Initialize an active connection.""" """Initialize an active connection."""
self.hass = hass self.hass = hass
self.request = request self.request = request
self.wsock: web.WebSocketResponse | None = None self.wsock = web.WebSocketResponse(heartbeat=55)
self._to_write: asyncio.Queue = asyncio.Queue(maxsize=MAX_PENDING_MSG) self._to_write: asyncio.Queue = asyncio.Queue(maxsize=MAX_PENDING_MSG)
self._handle_task: asyncio.Task | None = None self._handle_task: asyncio.Task | None = None
self._writer_task: asyncio.Task | None = None self._writer_task: asyncio.Task | None = None
@ -71,7 +71,6 @@ class WebSocketHandler:
async def _writer(self) -> None: async def _writer(self) -> None:
"""Write outgoing messages.""" """Write outgoing messages."""
# Exceptions if Socket disconnected or cancelled by connection handler # Exceptions if Socket disconnected or cancelled by connection handler
assert self.wsock is not None
with suppress(RuntimeError, ConnectionResetError, *CANCELLATION_ERRORS): with suppress(RuntimeError, ConnectionResetError, *CANCELLATION_ERRORS):
while not self.wsock.closed: while not self.wsock.closed:
message = await self._to_write.get() message = await self._to_write.get()
@ -143,8 +142,14 @@ class WebSocketHandler:
async def async_handle(self) -> web.WebSocketResponse: async def async_handle(self) -> web.WebSocketResponse:
"""Handle a websocket response.""" """Handle a websocket response."""
request = self.request request = self.request
wsock = self.wsock = web.WebSocketResponse(heartbeat=55) wsock = self.wsock
try:
async with async_timeout.timeout(10):
await wsock.prepare(request) await wsock.prepare(request)
except asyncio.TimeoutError:
self._logger.warning("Timeout preparing request from %s", request.remote)
return wsock
self._logger.debug("Connected from %s", request.remote) self._logger.debug("Connected from %s", request.remote)
self._handle_task = asyncio.current_task() self._handle_task = asyncio.current_task()

14
tests/components/websocket_api/test_http.py
View file

@ -1,8 +1,9 @@
"""Test Websocket API http module.""" """Test Websocket API http module."""
import asyncio
from datetime import timedelta from datetime import timedelta
from unittest.mock import patch from unittest.mock import patch
from aiohttp import WSMsgType from aiohttp import ServerDisconnectedError, WSMsgType, web
import pytest import pytest
from homeassistant.components.websocket_api import const, http from homeassistant.components.websocket_api import const, http
@ -80,3 +81,14 @@ async def test_non_json_message(hass, websocket_client, caplog):
f"Unable to serialize to JSON. Bad data found at $.result[0](State: test_domain.entity).attributes.bad={bad_data}(<class 'object'>" f"Unable to serialize to JSON. Bad data found at $.result[0](State: test_domain.entity).attributes.bad={bad_data}(<class 'object'>"
in caplog.text in caplog.text
) )
async def test_prepare_fail(hass, hass_ws_client, caplog):
"""Test failing to prepare."""
with patch(
"homeassistant.components.websocket_api.http.web.WebSocketResponse.prepare",
side_effect=(asyncio.TimeoutError, web.WebSocketResponse.prepare),
), pytest.raises(ServerDisconnectedError):
await hass_ws_client(hass)
assert "Timeout preparing request" in caplog.text