Adding stream cancel test

Author: blast-hardcheese

tests/v2/test_v2_cancellation.py

@@ -4,10 +4,12 @@
 from typing import (
     Any,
     AsyncIterator,
+    Literal,
 )
 
 import msgpack
 import nanoid
+import pytest
 
 from replit_river.messages import parse_transport_msg
 from replit_river.rpc import (
@@ -115,6 +117,151 @@ async def handle_server_messages() -> None:
     await server_handler
 
 
+@pytest.mark.parametrize("direction", ["send", "receive"])
+async def test_stream_cancel(
+    ws_server: WsServerFixture, direction: Literal["send", "receive"]
+) -> None:
+    (urimeta, recv, conn) = ws_server
+
+    client = Client(
+        client_id="CLIENT1",
+        server_id="SERVER",
+        transport_options=TransportOptions(),
+        uri_and_metadata_factory=urimeta,
+    )
+
+    connecting = asyncio.create_task(client.ensure_connected())
+    request_msg = parse_transport_msg(await recv.get())
+
+    assert not isinstance(request_msg, str)
+    assert (serverconn := conn())
+    handshake_request: ControlMessageHandshakeRequest[None] = (
+        ControlMessageHandshakeRequest(**request_msg.payload)
+    )
+
+    handshake_resp = ControlMessageHandshakeResponse(
+        status=HandShakeStatus(
+            ok=True,
+        ),
+    )
+    handshake_request.sessionId
+
+    msg = TransportMessage(
+        from_=request_msg.from_,
+        to=request_msg.to,
+        streamId=request_msg.streamId,
+        controlFlags=0,
+        id=nanoid.generate(),
+        seq=0,
+        ack=0,
+        payload=handshake_resp.model_dump(),
+    )
+    packed = msgpack.packb(
+        msg.model_dump(by_alias=True, exclude_none=True), datetime=True
+    )
+    await serverconn.send(packed)
+
+    bidi_waiter = asyncio.Event()
+
+    async def send_server_messages(request_msg: TransportMessage) -> None:
+        seq = 0
+
+        while True:
+            msg = TransportMessage(
+                from_=request_msg.to,
+                to=request_msg.from_,
+                streamId=request_msg.streamId,
+                controlFlags=0,
+                id=nanoid.generate(),
+                seq=seq,
+                ack=0,
+                payload={
+                    "ok": True,
+                    "payload": {
+                        "hello": "world",
+                    },
+                },
+            )
+            seq += 1
+            packed = msgpack.packb(
+                msg.model_dump(by_alias=True, exclude_none=True), datetime=True
+            )
+            await serverconn.send(packed)
+            await asyncio.sleep(0.1)
+
+            if seq > 5 and direction == "send":
+                bidi_waiter.set()
+
+    async def handle_server_messages(request_msg: TransportMessage) -> None:
+        msg = TransportMessage(**msgpack.unpackb(await recv.get()))
+        while msg.payload.get("payload", {}).get("hello") == "world":
+            logging.debug("Found a hello:world %r", repr(msg))
+            msg = TransportMessage(**msgpack.unpackb(await recv.get()))
+
+        assert msg.controlFlags == STREAM_CANCEL_BIT
+
+    async def receive_chunks() -> None:
+        async def _upload_chunks() -> AsyncIterator[OuterPayload[dict[Any, Any]]]:
+            count = 0
+            while True:
+                await asyncio.sleep(0.1)
+                yield {
+                    "ok": True,
+                    "payload": {
+                        "hello": "world",
+                    },
+                }
+                count += 1
+                if count > 5 and direction == "receive":
+                    # We've sent enough messages, interrupt the stream.
+                    bidi_waiter.set()
+
+        async for chunk in client.send_stream(
+            "test",
+            "bigstream",
+            {},
+            _upload_chunks(),
+            lambda x: x,
+            lambda x: x,
+            lambda x: x,
+            lambda x: x,
+        ):
+            print(repr(chunk))
+
+    receive_task = asyncio.create_task(receive_chunks())
+    request_msg = parse_transport_msg(await recv.get())
+    logging.debug("request_msg: %r", repr(request_msg))
+    assert not isinstance(request_msg, str)
+
+    server_sender = asyncio.create_task(send_server_messages(request_msg))
+    server_receiver = asyncio.create_task(handle_server_messages(request_msg))
+
+    # Wait until we've seen at least a few messages from the requisite Task
+    await bidi_waiter.wait()
+
+    receive_task.cancel()
+    try:
+        await receive_task
+    except asyncio.CancelledError:
+        pass
+
+    await client.close()
+    await connecting
+
+    # Ensure we're listening to close messages as well
+    assert server_sender
+    server_sender.cancel()
+    try:
+        await server_sender
+    except asyncio.CancelledError:
+        pass
+    server_receiver.cancel()
+    try:
+        await server_receiver
+    except Exception:
+        pass
+
+
 async def test_subscription_cancel(ws_server: WsServerFixture) -> None:
     (urimeta, recv, conn) = ws_server