[PR #9566/22f0831 backport][3.11] Refactor WebSocketWriter to remove high level protocol functions (#9569)

Author: bdraco

aiohttp/_websocket/writer.py

@@ -34,6 +34,14 @@
 
 
 class WebSocketWriter:
+    """WebSocket writer.
+
+    The writer is responsible for sending messages to the client. It is
+    created by the protocol when a connection is established. The writer
+    should avoid implementing any application logic and should only be
+    concerned with the low-level details of the WebSocket protocol.
+    """
+
     def __init__(
         self,
         protocol: BaseProtocol,
@@ -45,6 +53,7 @@ def __init__(
         compress: int = 0,
         notakeover: bool = False,
     ) -> None:
+        """Initialize a WebSocket writer."""
         self.protocol = protocol
         self.transport = transport
         self.use_mask = use_mask
@@ -155,18 +164,6 @@ def _make_compress_obj(self, compress: int) -> ZLibCompressor:
             max_sync_chunk_size=WEBSOCKET_MAX_SYNC_CHUNK_SIZE,
         )
 
-    async def pong(self, message: Union[bytes, str] = b"") -> None:
-        """Send pong message."""
-        if isinstance(message, str):
-            message = message.encode("utf-8")
-        await self.send_frame(message, WSMsgType.PONG)
-
-    async def ping(self, message: Union[bytes, str] = b"") -> None:
-        """Send ping message."""
-        if isinstance(message, str):
-            message = message.encode("utf-8")
-        await self.send_frame(message, WSMsgType.PING)
-
     async def close(self, code: int = 1000, message: Union[bytes, str] = b"") -> None:
         """Close the websocket, sending the specified code and message."""
         if isinstance(message, str):

aiohttp/client_ws.py

@@ -140,13 +140,14 @@ def _send_heartbeat(self) -> None:
         self._cancel_pong_response_cb()
         self._pong_response_cb = loop.call_at(when, self._pong_not_received)
 
+        coro = self._writer.send_frame(b"", WSMsgType.PING)
         if sys.version_info >= (3, 12):
             # Optimization for Python 3.12, try to send the ping
             # immediately to avoid having to schedule
             # the task on the event loop.
-            ping_task = asyncio.Task(self._writer.ping(), loop=loop, eager_start=True)
+            ping_task = asyncio.Task(coro, loop=loop, eager_start=True)
         else:
-            ping_task = loop.create_task(self._writer.ping())
+            ping_task = loop.create_task(coro)
 
         if not ping_task.done():
             self._ping_task = ping_task
@@ -224,10 +225,10 @@ def exception(self) -> Optional[BaseException]:
         return self._exception
 
     async def ping(self, message: bytes = b"") -> None:
-        await self._writer.ping(message)
+        await self._writer.send_frame(message, WSMsgType.PING)
 
     async def pong(self, message: bytes = b"") -> None:
-        await self._writer.pong(message)
+        await self._writer.send_frame(message, WSMsgType.PONG)
 
     async def send_frame(
         self, message: bytes, opcode: WSMsgType, compress: Optional[int] = None

aiohttp/web_ws.py

@@ -153,13 +153,14 @@ def _send_heartbeat(self) -> None:
         self._cancel_pong_response_cb()
         self._pong_response_cb = loop.call_at(when, self._pong_not_received)
 
+        coro = self._writer.send_frame(b"", WSMsgType.PING)
         if sys.version_info >= (3, 12):
             # Optimization for Python 3.12, try to send the ping
             # immediately to avoid having to schedule
             # the task on the event loop.
-            ping_task = asyncio.Task(self._writer.ping(), loop=loop, eager_start=True)
+            ping_task = asyncio.Task(coro, loop=loop, eager_start=True)
         else:
-            ping_task = loop.create_task(self._writer.ping())
+            ping_task = loop.create_task(coro)
 
         if not ping_task.done():
             self._ping_task = ping_task
@@ -371,13 +372,13 @@ def exception(self) -> Optional[BaseException]:
     async def ping(self, message: bytes = b"") -> None:
         if self._writer is None:
             raise RuntimeError("Call .prepare() first")
-        await self._writer.ping(message)
+        await self._writer.send_frame(message, WSMsgType.PING)
 
     async def pong(self, message: bytes = b"") -> None:
         # unsolicited pong
         if self._writer is None:
             raise RuntimeError("Call .prepare() first")
-        await self._writer.pong(message)
+        await self._writer.send_frame(message, WSMsgType.PONG)
 
     async def send_frame(
         self, message: bytes, opcode: WSMsgType, compress: Optional[int] = None

tests/test_client_ws_functional.py

@@ -1,6 +1,6 @@
 import asyncio
 import sys
-from typing import Any, NoReturn
+from typing import Any, NoReturn, Optional
 from unittest import mock
 
 import pytest
@@ -704,9 +704,11 @@ async def handler(request: web.Request) -> NoReturn:
     assert resp._conn is not None
     with mock.patch.object(
         resp._conn.transport, "write", side_effect=ClientConnectionResetError
-    ), mock.patch.object(resp._writer, "ping", wraps=resp._writer.ping) as ping:
+    ), mock.patch.object(
+        resp._writer, "send_frame", wraps=resp._writer.send_frame
+    ) as send_frame:
         await resp.receive()
-        ping_count = ping.call_count
+        ping_count = send_frame.call_args_list.count(mock.call(b"", WSMsgType.PING))
     # Connection should be closed roughly after 1.5x heartbeat.
     await asyncio.sleep(0.2)
     assert ping_count == 1
@@ -842,7 +844,7 @@ async def handler(request):
 
 
 async def test_close_websocket_while_ping_inflight(
-    aiohttp_client: AiohttpClient,
+    aiohttp_client: AiohttpClient, loop: asyncio.AbstractEventLoop
 ) -> None:
     """Test closing the websocket while a ping is in-flight."""
     ping_received = False
@@ -866,23 +868,27 @@ async def handler(request: web.Request) -> NoReturn:
     await resp.send_bytes(b"ask")
 
     cancelled = False
-    ping_stated = False
-
-    async def delayed_ping() -> None:
-        nonlocal cancelled, ping_stated
-        ping_stated = True
+    ping_started = loop.create_future()
+
+    async def delayed_send_frame(
+        message: bytes, opcode: int, compress: Optional[int] = None
+    ) -> None:
+        assert opcode == WSMsgType.PING
+        nonlocal cancelled, ping_started
+        ping_started.set_result(None)
         try:
             await asyncio.sleep(1)
         except asyncio.CancelledError:
             cancelled = True
             raise
 
-    with mock.patch.object(resp._writer, "ping", delayed_ping):
-        await asyncio.sleep(0.1)
+    with mock.patch.object(resp._writer, "send_frame", delayed_send_frame):
+        async with async_timeout.timeout(1):
+            await ping_started
 
     await resp.close()
     await asyncio.sleep(0)
-    assert ping_stated is True
+    assert ping_started.result() is None
     assert cancelled is True
 
 

tests/test_web_websocket_functional.py

@@ -423,7 +423,7 @@ async def handler(request):
         ws = web.WebSocketResponse()
         await ws.prepare(request)
 
-        await ws.ping("data")
+        await ws.ping(b"data")
         await ws.receive()
         closed.set_result(None)
         return ws
@@ -460,7 +460,7 @@ async def handler(request):
 
     ws = await client.ws_connect("/", autoping=False)
 
-    await ws.ping("data")
+    await ws.ping(b"data")
     msg = await ws.receive()
     assert msg.type == WSMsgType.PONG
     assert msg.data == b"data"
@@ -478,7 +478,7 @@ async def handler(request):
 
         msg = await ws.receive()
         assert msg.type == WSMsgType.PING
-        await ws.pong("data")
+        await ws.pong(b"data")
 
         msg = await ws.receive()
         assert msg.type == WSMsgType.CLOSE
@@ -493,7 +493,7 @@ async def handler(request):
 
     ws = await client.ws_connect("/", autoping=False)
 
-    await ws.ping("data")
+    await ws.ping(b"data")
     msg = await ws.receive()
     assert msg.type == WSMsgType.PONG
     assert msg.data == b"data"
@@ -741,12 +741,14 @@ async def handler(request: web.Request) -> NoReturn:
         with mock.patch.object(
             ws_server._req.transport, "write", side_effect=ConnectionResetError
         ), mock.patch.object(
-            ws_server._writer, "ping", wraps=ws_server._writer.ping
-        ) as ping:
+            ws_server._writer, "send_frame", wraps=ws_server._writer.send_frame
+        ) as send_frame:
             try:
                 await ws_server.receive()
             finally:
-                ping_count = ping.call_count
+                ping_count = send_frame.call_args_list.count(
+                    mock.call(b"", WSMsgType.PING)
+                )
         assert False
 
     app = web.Application()
@@ -990,7 +992,7 @@ async def handler(request):
         msg = await ws.receive()
         assert msg.type == WSMsgType.PING
         await asyncio.sleep(0)
-        await ws.pong("data")
+        await ws.pong(b"data")
 
         msg = await ws.receive()
         assert msg.type == WSMsgType.CLOSE
@@ -1006,7 +1008,7 @@ async def handler(request):
     ws = await client.ws_connect("/", autoping=False)
 
     await asyncio.sleep(0)
-    await ws.ping("data")
+    await ws.ping(b"data")
 
     msg = await ws.receive()
     assert msg.type == WSMsgType.PONG
@@ -1036,7 +1038,7 @@ async def handler(request):
         msg = await ws.receive()
         assert msg.type == WSMsgType.PING
         await asyncio.sleep(0)
-        await ws.pong("data")
+        await ws.pong(b"data")
 
         msg = await ws.receive()
         assert msg.type == WSMsgType.CLOSE
@@ -1052,7 +1054,7 @@ async def handler(request):
     ws = await client.ws_connect("/", autoping=False)
 
     await timed_out
-    await ws.ping("data")
+    await ws.ping(b"data")
 
     msg = await ws.receive()
     assert msg.type == WSMsgType.PONG

tests/test_websocket_writer.py

@@ -29,14 +29,14 @@ def writer(protocol, transport):
     return WebSocketWriter(protocol, transport, use_mask=False)
 
 
-async def test_pong(writer) -> None:
-    await writer.pong()
-    writer.transport.write.assert_called_with(b"\x8a\x00")
+async def test_pong(writer: WebSocketWriter) -> None:
+    await writer.send_frame(b"", WSMsgType.PONG)
+    writer.transport.write.assert_called_with(b"\x8a\x00")  # type: ignore[attr-defined]
 
 
-async def test_ping(writer) -> None:
-    await writer.ping()
-    writer.transport.write.assert_called_with(b"\x89\x00")
+async def test_ping(writer: WebSocketWriter) -> None:
+    await writer.send_frame(b"", WSMsgType.PING)
+    writer.transport.write.assert_called_with(b"\x89\x00")  # type: ignore[attr-defined]
 
 
 async def test_send_text(writer: WebSocketWriter) -> None: