Adding an RPC cancellation test

Author: blast-hardcheese

tests/v2/test_v2_cancellation.py

@@ -1,5 +1,6 @@
 import asyncio
 import logging
+from datetime import timedelta
 from typing import (
     Any,
     AsyncIterator,
@@ -10,17 +11,110 @@
 
 from replit_river.messages import parse_transport_msg
 from replit_river.rpc import (
+    STREAM_OPEN_BIT,
     ControlMessageHandshakeRequest,
     ControlMessageHandshakeResponse,
     HandShakeStatus,
     TransportMessage,
 )
 from replit_river.transport_options import TransportOptions
 from replit_river.v2.client import Client
-from replit_river.v2.session import STREAM_CANCEL_BIT
+from replit_river.v2.session import STREAM_CANCEL_BIT, STREAM_CLOSED_BIT
 from tests.v2.fixtures.raw_ws_server import OuterPayload, WsServerFixture
 
 
+async def test_rpc_cancel(ws_server: WsServerFixture) -> 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)
+
+    sent_waiter = asyncio.Event()
+
+    async def handle_server_messages() -> None:
+        request_msg = parse_transport_msg(await recv.get())
+        assert not isinstance(request_msg, str)
+
+        logging.debug("request_msg: %r", repr(request_msg))
+
+        assert request_msg.payload.get("payload", {}).get("hello") == "world"
+        logging.debug("Found a hello:world %r", repr(request_msg))
+
+        sent_waiter.set()
+
+        assert request_msg.controlFlags == STREAM_OPEN_BIT | STREAM_CLOSED_BIT
+
+        cancel_msg = parse_transport_msg(await recv.get())
+        assert not isinstance(cancel_msg, str)
+        assert cancel_msg.controlFlags == STREAM_CANCEL_BIT
+
+    server_handler = asyncio.create_task(handle_server_messages())
+
+    rpc_task = asyncio.create_task(
+        client.send_rpc(
+            "test",
+            "bigstream",
+            {"ok": True, "payload": {"hello": "world"}},
+            lambda x: x,
+            lambda x: x,
+            lambda x: x,
+            timedelta(seconds=2),
+        )
+    )
+
+    # Wait until we've seen at least a few messages from the upload Task
+    await sent_waiter.wait()
+
+    rpc_task.cancel()
+
+    try:
+        await rpc_task
+    except asyncio.CancelledError:
+        pass
+
+    await client.close()
+    await connecting
+
+    # Ensure we're listening to close messages as well
+    server_handler.cancel()
+    await server_handler
+
+
 async def test_upload_cancel(ws_server: WsServerFixture) -> None:
     (urimeta, recv, conn) = ws_server