Leave websocket transport open if receive times out or is cancelled (#…

…8251)

(cherry picked from commit c21b76d)

CHANGES/8251.bugfix.rst

@@ -0,0 +1,4 @@
+Leave websocket transport open if receive times out or is cancelled
+-- by :user:`bdraco`.
+
+This restores the behavior prior to the change in #7978.

aiohttp/web_ws.py

@@ -462,8 +462,7 @@ async def receive(self, timeout: Optional[float] = None) -> WSMessage:
  waiter = self._waiting
  set_result(waiter, True)
  self._waiting = None
- except (asyncio.CancelledError, asyncio.TimeoutError):
- self._set_code_close_transport(WSCloseCode.ABNORMAL_CLOSURE)
+ except asyncio.TimeoutError:
  raise
  except EofStream:
  self._close_code = WSCloseCode.OK

tests/test_web_websocket.py

@@ -472,7 +472,8 @@ async def test_receive_timeouterror(make_request: Any, loop: Any) -> None:
  with pytest.raises(asyncio.TimeoutError):
  await ws.receive()
 
- assert len(ws._req.transport.close.mock_calls) == 1
+ # Should not close the connection on timeout
+ assert len(ws._req.transport.close.mock_calls) == 0
 
 
 async def test_multiple_receive_on_close_connection(make_request) -> None:

tests/test_web_websocket_functional.py

@@ -1,6 +1,8 @@
 # HTTP websocket server functional tests
 
 import asyncio
+import contextlib
+import sys
 from typing import Any, Optional
 
 import pytest
@@ -797,3 +799,94 @@ async def ws_handler(request):
  resp = await client.get("/api/null", timeout=1)
  assert (await resp.json()) == {"err": None}
  resp.close()
+
+
+async def test_receive_being_cancelled_keeps_connection_open(
+ loop: Any, aiohttp_client: Any
+) -> None:
+ closed = loop.create_future()
+
+ async def handler(request):
+ ws = web.WebSocketResponse(autoping=False)
+ await ws.prepare(request)
+
+ task = asyncio.create_task(ws.receive())
+ await asyncio.sleep(0)
+ task.cancel()
+ with contextlib.suppress(asyncio.CancelledError):
+ await task
+
+ msg = await ws.receive()
+ assert msg.type == WSMsgType.PING
+ await asyncio.sleep(0)
+ await ws.pong("data")
+
+ msg = await ws.receive()
+ assert msg.type == WSMsgType.CLOSE
+ assert msg.data == WSCloseCode.OK
+ assert msg.extra == "exit message"
+ closed.set_result(None)
+ return ws
+
+ app = web.Application()
+ app.router.add_get("/", handler)
+ client = await aiohttp_client(app)
+
+ ws = await client.ws_connect("/", autoping=False)
+
+ await asyncio.sleep(0)
+ await ws.ping("data")
+
+ msg = await ws.receive()
+ assert msg.type == WSMsgType.PONG
+ assert msg.data == b"data"
+
+ await ws.close(code=WSCloseCode.OK, message="exit message")
+
+ await closed
+
+
+async def test_receive_timeout_keeps_connection_open(
+ loop: Any, aiohttp_client: Any
+) -> None:
+ closed = loop.create_future()
+ timed_out = loop.create_future()
+
+ async def handler(request):
+ ws = web.WebSocketResponse(autoping=False)
+ await ws.prepare(request)
+
+ task = asyncio.create_task(ws.receive(sys.float_info.min))
+ with contextlib.suppress(asyncio.TimeoutError):
+ await task
+
+ timed_out.set_result(None)
+
+ msg = await ws.receive()
+ assert msg.type == WSMsgType.PING
+ await asyncio.sleep(0)
+ await ws.pong("data")
+
+ msg = await ws.receive()
+ assert msg.type == WSMsgType.CLOSE
+ assert msg.data == WSCloseCode.OK
+ assert msg.extra == "exit message"
+ closed.set_result(None)
+ return ws
+
+ app = web.Application()
+ app.router.add_get("/", handler)
+ client = await aiohttp_client(app)
+
+ ws = await client.ws_connect("/", autoping=False)
+
+ await timed_out
+ await ws.ping("data")
+
+ msg = await ws.receive()
+ assert msg.type == WSMsgType.PONG
+ assert msg.data == b"data"
+
+ await ws.close(code=WSCloseCode.OK, message="exit message")
+
+ await closed