Adding a test for cancelling upload

Author: blast-hardcheese

tests/v2/test_v2_session_lifecycle.py

@@ -1,6 +1,14 @@
 import asyncio
 import logging
-from typing import AsyncIterator, Awaitable, Callable, TypeAlias, TypedDict
+from typing import (
+    Any,
+    AsyncIterator,
+    Awaitable,
+    Callable,
+    Literal,
+    TypeAlias,
+    TypedDict,
+)
 
 import msgpack
 import nanoid
@@ -20,7 +28,12 @@
 )
 from replit_river.transport_options import TransportOptions, UriAndMetadata
 from replit_river.v2.client import Client
-from replit_river.v2.session import STREAM_CLOSED_BIT, Session
+from replit_river.v2.session import STREAM_CANCEL_BIT, STREAM_CLOSED_BIT, Session
+
+
+class OuterPayload[A](TypedDict):
+    ok: Literal[True]
+    payload: A
 
 
 class _PermissiveRateLimiter(RateLimiter):
@@ -247,3 +260,104 @@ async def handle_server_messages() -> None:
     # Ensure we're listening to close messages as well
     stream_handler.cancel()
     await stream_handler
+
+
+async def test_upload_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)
+
+    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))
+
+        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
+
+    stream_handler = asyncio.create_task(handle_server_messages())
+
+    sent_waiter = asyncio.Event()
+
+    async def upload_chunks() -> AsyncIterator[OuterPayload[dict[Any, Any]]]:
+        count = 0
+        while True:
+            yield {
+                "ok": True,
+                "payload": {
+                    "hello": "world",
+                },
+            }
+            count += 1
+            if count > 5:
+                sent_waiter.set()
+                await asyncio.sleep(0.1)
+
+    upload_task = asyncio.create_task(
+        client.send_upload(
+            "test",
+            "bigstream",
+            {},
+            upload_chunks(),
+            lambda x: x,
+            lambda x: x,
+            lambda x: x,
+            lambda x: x,
+        )
+    )
+
+    await sent_waiter.wait()
+
+    upload_task.cancel()
+    try:
+        await upload_task
+    except asyncio.CancelledError:
+        pass
+
+    await client.close()
+    await connecting
+
+    # Ensure we're listening to close messages as well
+    stream_handler.cancel()
+    await stream_handler