chuanqi/cq2_sdk/tests/unit/test_tcp_client.py

import asyncio
import pytest
import cq2_sdk.transport.tcp as tcp_module
from cq2_sdk.transport.tcp import TcpClient
from cq2_sdk.transport.framing import WIRE_HEADER_SIZE, decode_header, encode_frame


async def _echo_handler(reader, writer):
    try:
        while True:
            data = await reader.read(4096)
            if not data:
                break
            writer.write(data)
            await writer.drain()
    finally:
        writer.close()


@pytest.fixture
async def echo_server():
    server = await asyncio.start_server(_echo_handler, "127.0.0.1", 0)
    port = server.sockets[0].getsockname()[1]
    yield port
    server.close()
    await server.wait_closed()


@pytest.mark.asyncio
async def test_send_recv_roundtrip(echo_server):
    cli = TcpClient("127.0.0.1", echo_server)
    await cli.connect()
    try:
        await cli.send(msg_id=2001, recog=900000, param1=1, param2=0, param3=0, body=b"hello")
        hdr, body = await cli.recv(timeout=2.0)
        assert hdr.msg_id == 2001
        assert hdr.recog == 900000
        assert hdr.param1 == 1
        assert body == b"hello"
    finally:
        await cli.close()


@pytest.mark.asyncio
async def test_send_multiple_frames(echo_server):
    cli = TcpClient("127.0.0.1", echo_server)
    await cli.connect()
    try:
        for i in range(3):
            await cli.send(msg_id=100 + i, recog=0, param1=i, param2=0, param3=0, body=f"frame{i}".encode())
        for i in range(3):
            hdr, body = await cli.recv(timeout=2.0)
            assert hdr.msg_id == 100 + i
            assert body == f"frame{i}".encode()
    finally:
        await cli.close()


@pytest.mark.asyncio
async def test_recv_timeout_raises():
    cli = TcpClient("127.0.0.1", 1)  # nothing listening
    with pytest.raises(asyncio.TimeoutError):
        await cli.recv(timeout=0.3)


@pytest.mark.asyncio
async def test_close_is_idempotent(echo_server):
    cli = TcpClient("127.0.0.1", echo_server)
    await cli.connect()
    await cli.close()
    await cli.close()  # should not raise
    assert cli._closed


@pytest.mark.asyncio
async def test_send_before_connect_raises():
    cli = TcpClient("127.0.0.1", 1)
    with pytest.raises(RuntimeError, match="not connected"):
        await cli.send(1, 0, 0, 0, 0, b"")


@pytest.mark.asyncio
async def test_recv_reassembles_split_wire_header_and_body():
    async def handler(reader, writer):
        wire = encode_frame(529, 1, 0, 0, 0, b"split", sequence=17)
        for chunk in (wire[:11], wire[11:WIRE_HEADER_SIZE], wire[WIRE_HEADER_SIZE:]):
            writer.write(chunk)
            await writer.drain()
        writer.close()

    server = await asyncio.start_server(handler, "127.0.0.1", 0)
    port = server.sockets[0].getsockname()[1]
    client = TcpClient("127.0.0.1", port)
    await client.connect()
    try:
        header, body = await client.recv(timeout=1.0)
        assert header.msg_id == 529
        assert header.sequence == 17
        assert body == b"split"
    finally:
        await client.close()
        server.close()
        await server.wait_closed()


@pytest.mark.asyncio
async def test_send_uses_and_increments_initial_sequence():
    received: list[int] = []
    complete = asyncio.Event()

    async def handler(reader, writer):
        for _ in range(2):
            header = decode_header(await reader.readexactly(WIRE_HEADER_SIZE))
            if header.body_len:
                await reader.readexactly(header.body_len)
            received.append(header.sequence)
        complete.set()
        writer.close()

    server = await asyncio.start_server(handler, "127.0.0.1", 0)
    port = server.sockets[0].getsockname()[1]
    client = TcpClient("127.0.0.1", port, initial_sequence=9)
    await client.connect()
    try:
        await client.send(2001, 0, 0, 0, 0, b"first")
        await client.send(2002, 0, 0, 0, 0, b"second")
        await asyncio.wait_for(complete.wait(), timeout=1.0)
        assert received == [9, 10]
    finally:
        await client.close()
        server.close()
        await server.wait_closed()


@pytest.mark.asyncio
async def test_send_uses_current_sdata():
    received = asyncio.Future()

    async def handler(reader, writer):
        header = decode_header(await reader.readexactly(WIRE_HEADER_SIZE))
        if header.body_len:
            await reader.readexactly(header.body_len)
        received.set_result(header)
        writer.close()

    server = await asyncio.start_server(handler, "127.0.0.1", 0)
    port = server.sockets[0].getsockname()[1]
    client = TcpClient("127.0.0.1", port)
    client.set_sdata(149)
    await client.connect()
    try:
        await client.send(2001, 900000, 1, 2, 3, b"login")
        header = await asyncio.wait_for(received, timeout=1.0)
        assert header.param3 == 3
        assert header.sdata == 149
    finally:
        await client.close()
        server.close()
        await server.wait_closed()


@pytest.mark.asyncio
async def test_reconnect_resets_sequence_and_preserves_sdata():
    received: list[tuple[int, int]] = []
    complete = asyncio.Event()

    async def handler(reader, writer):
        header = decode_header(await reader.readexactly(WIRE_HEADER_SIZE))
        if header.body_len:
            await reader.readexactly(header.body_len)
        received.append((header.sequence, header.sdata))
        if len(received) == 2:
            complete.set()
        writer.close()
        await writer.wait_closed()

    server = await asyncio.start_server(handler, "127.0.0.1", 0)
    port = server.sockets[0].getsockname()[1]
    client = TcpClient("127.0.0.1", port)
    client.set_sdata(149)
    await client.connect()
    try:
        await client.send(100, 0, 0, 0, 0, b"role")
        await client.reconnect("127.0.0.1", port)
        await client.send(20, 0, 0, 0, 0, b"world")
        await asyncio.wait_for(complete.wait(), timeout=1.0)
        assert received == [(0, 149), (0, 149)]
    finally:
        await client.close()
        server.close()
        await server.wait_closed()


@pytest.mark.asyncio
async def test_heartbeat_uses_current_unix_time_in_recog(monkeypatch):
    received = asyncio.Future()

    async def handler(reader, writer):
        header = decode_header(await reader.readexactly(WIRE_HEADER_SIZE))
        received.set_result(header)
        writer.close()

    class FixedTime:
        @staticmethod
        def time():
            return 1_789_709_744

    monkeypatch.setattr(tcp_module, "time", FixedTime, raising=False)
    server = await asyncio.start_server(handler, "127.0.0.1", 0)
    port = server.sockets[0].getsockname()[1]
    client = TcpClient("127.0.0.1", port)
    client.set_sdata(149)
    await client.connect()
    try:
        await client.start_heartbeat(interval=0.01)
        header = await asyncio.wait_for(received, timeout=1.0)
        assert header.msg_id == 60000
        assert header.recog == 1_789_709_744
        assert header.sdata == 149
    finally:
        await client.close()
        server.close()
        await server.wait_closed()


@pytest.mark.asyncio
async def test_recv_surfaces_unsupported_body_codec():
    async def handler(reader, writer):
        writer.write(encode_frame(529, 0, 0, 0, 0, b"compressed", codec=1))
        await writer.drain()
        writer.close()

    server = await asyncio.start_server(handler, "127.0.0.1", 0)
    port = server.sockets[0].getsockname()[1]
    client = TcpClient("127.0.0.1", port)
    await client.connect()
    try:
        with pytest.raises(ValueError, match="unsupported body codec: 1"):
            await client.recv(timeout=1.0)
    finally:
        await client.close()
        server.close()
        await server.wait_closed()


@pytest.mark.asyncio
async def test_recv_delivers_queued_frame_before_later_codec_error():
    async def handler(reader, writer):
        writer.write(encode_frame(529, 0, 0, 0, 0, b"valid"))
        writer.write(encode_frame(530, 0, 0, 0, 0, b"compressed", codec=1))
        await writer.drain()
        writer.close()

    server = await asyncio.start_server(handler, "127.0.0.1", 0)
    port = server.sockets[0].getsockname()[1]
    client = TcpClient("127.0.0.1", port)
    await client.connect()
    try:
        for _ in range(100):
            if client._recv_error is not None:
                break
            await asyncio.sleep(0.01)
        assert client._recv_error is not None
        header, body = await client.recv(timeout=1.0)
        assert header.msg_id == 529
        assert body == b"valid"
        with pytest.raises(ValueError, match="unsupported body codec: 1"):
            await client.recv(timeout=1.0)
    finally:
        await client.close()
        server.close()
        await server.wait_closed()