[PR #9698/422ba0c backport][3.11] Add flow control tests for WebSocket (#9701)

Co-authored-by: pre-commit-ci[bot]

Author: bdraco

tests/test_websocket_parser.py

@@ -18,6 +18,7 @@
     websocket_mask,
 )
 from aiohttp._websocket.models import WS_DEFLATE_TRAILING
+from aiohttp.base_protocol import BaseProtocol
 from aiohttp.http import WebSocketError, WSCloseCode, WSMessage, WSMsgType
 from aiohttp.http_websocket import WebSocketReader
 
@@ -83,12 +84,36 @@ def build_close_frame(code=1000, message=b"", noheader=False):
 
 
 @pytest.fixture()
-def out(loop):
-    return aiohttp.DataQueue(loop)
+def protocol(loop: asyncio.AbstractEventLoop) -> BaseProtocol:
+    transport = mock.Mock(spec_set=asyncio.Transport)
+    protocol = BaseProtocol(loop)
+    protocol.connection_made(transport)
+    return protocol
 
 
 @pytest.fixture()
-def parser(out: aiohttp.DataQueue[WSMessage]) -> PatchableWebSocketReader:
+def out(
+    loop: asyncio.AbstractEventLoop, protocol: BaseProtocol
+) -> aiohttp.FlowControlDataQueue[WSMessage]:
+    return aiohttp.FlowControlDataQueue(protocol, 2**16, loop=loop)
+
+
+@pytest.fixture()
+def out_low_limit(
+    loop: asyncio.AbstractEventLoop, protocol: BaseProtocol
+) -> aiohttp.FlowControlDataQueue[WSMessage]:
+    return aiohttp.FlowControlDataQueue(protocol, 16, loop=loop)
+
+
+@pytest.fixture()
+def parser_low_limit(
+    out_low_limit: aiohttp.FlowControlDataQueue[WSMessage],
+) -> PatchableWebSocketReader:
+    return PatchableWebSocketReader(out_low_limit, 4 * 1024 * 1024)
+
+
+@pytest.fixture()
+def parser(out: aiohttp.FlowControlDataQueue[WSMessage]) -> PatchableWebSocketReader:
     return PatchableWebSocketReader(out, 4 * 1024 * 1024)
 
 
@@ -569,3 +594,42 @@ def test_pickle(self) -> None:
             assert err2.code == WSCloseCode.PROTOCOL_ERROR
             assert str(err2) == "Something invalid"
             assert err2.foo == "bar"
+
+
+def test_flow_control_binary(
+    protocol: BaseProtocol,
+    out_low_limit: aiohttp.FlowControlDataQueue[WSMessage],
+    parser_low_limit: WebSocketReader,
+) -> None:
+    large_payload = b"b" * (1 + 16 * 2)
+    large_payload_len = len(large_payload)
+    with mock.patch.object(parser_low_limit, "parse_frame", autospec=True) as m:
+        m.return_value = [(1, WSMsgType.BINARY, large_payload, False)]
+
+        parser_low_limit.feed_data(b"")
+
+    res = out_low_limit._buffer[0]
+    assert res == (WSMessage(WSMsgType.BINARY, large_payload, ""), large_payload_len)
+    assert protocol._reading_paused is True
+
+
+@pytest.mark.xfail(
+    reason="Flow control is currently broken on master branch; see #9686"
+)
+def test_flow_control_multi_byte_text(
+    protocol: BaseProtocol,
+    out_low_limit: aiohttp.FlowControlDataQueue[WSMessage],
+    parser_low_limit: WebSocketReader,
+) -> None:
+    large_payload_text = "𒀁" * (1 + 16 * 2)
+    large_payload = large_payload_text.encode("utf-8")
+    large_payload_len = len(large_payload)
+
+    with mock.patch.object(parser_low_limit, "parse_frame", autospec=True) as m:
+        m.return_value = [(1, WSMsgType.TEXT, large_payload, False)]
+
+        parser_low_limit.feed_data(b"")
+
+    res = out_low_limit._buffer[0]
+    assert res == (WSMessage(WSMsgType.TEXT, large_payload_text, ""), large_payload_len)
+    assert protocol._reading_paused is True