Add a subscription test

Author: blast-hardcheese

tests/v2/test_v2_cancellation.py

@@ -115,6 +115,123 @@ async def handle_server_messages() -> None:
     await server_handler
 
 
+async def test_subscription_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)
+
+    received_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))
+        seq = 0
+
+        while True:
+            try:
+                cancel_msg = parse_transport_msg(recv.get_nowait())
+                break
+            except asyncio.queues.QueueEmpty:
+                pass
+
+            msg = TransportMessage(
+                from_=request_msg.from_,
+                to=request_msg.to,
+                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:
+                received_waiter.set()
+
+        assert not isinstance(cancel_msg, str)
+        assert cancel_msg.controlFlags == STREAM_CANCEL_BIT
+
+    server_handler = asyncio.create_task(handle_server_messages())
+
+    async def receive_chunks() -> None:
+        async for chunk in client.send_subscription(
+            "test",
+            "bigstream",
+            {},
+            lambda x: x,
+            lambda x: x,
+            lambda x: x,
+        ):
+            print(repr(chunk))
+
+    receive_task = asyncio.create_task(receive_chunks())
+
+    # Wait until we've seen at least a few messages from the upload Task
+    await received_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
+    server_handler.cancel()
+    await server_handler
+
+
 async def test_upload_cancel(ws_server: WsServerFixture) -> None:
     (urimeta, recv, conn) = ws_server