proto-socket/python/test/test_communicator.py

from __future__ import annotations

import asyncio

import pytest

from proto_socket.communicator import MAX_NONCE, Communicator, type_name_of
from proto_socket.packets.message_common_pb2 import PacketBase, TestData as ProtoTestData


class FakeTransport:
    def __init__(self) -> None:
        self.packets: list[PacketBase] = []
        self.closed = False

    async def write_packet(self, base: PacketBase) -> None:
        self.packets.append(base)

    async def close(self) -> None:
        self.closed = True


def parser_map():
    return {type_name_of(ProtoTestData): ProtoTestData.FromString}


@pytest.mark.asyncio
async def test_add_listener_exclusivity():
    transport = FakeTransport()
    communicator = Communicator()
    communicator.initialize(transport, parser_map())
    communicator.add_listener(type_name_of(ProtoTestData), lambda _: None)

    with pytest.raises(ValueError):
        communicator.add_request_listener(type_name_of(ProtoTestData), lambda req: req)

    await communicator.close()


@pytest.mark.asyncio
async def test_send_and_receive():
    transport = FakeTransport()
    communicator = Communicator()
    communicator.initialize(transport, parser_map())
    received = asyncio.get_running_loop().create_future()
    communicator.add_listener(type_name_of(ProtoTestData), received.set_result)

    await communicator.send(ProtoTestData(index=7, message="hello"))
    assert len(transport.packets) == 1
    packet = transport.packets[0]
    assert packet.typeName == "TestData"
    assert packet.nonce == 1

    communicator.on_received_data(packet.typeName, packet.data, packet.nonce, 0)
    result = await asyncio.wait_for(received, 1)
    assert result.index == 7
    assert result.message == "hello"

    await communicator.close()


@pytest.mark.asyncio
async def test_send_request_response():
    transport = FakeTransport()
    communicator = Communicator()
    communicator.initialize(transport, parser_map())

    task = asyncio.create_task(
        communicator.send_request(ProtoTestData(index=5, message="request"), ProtoTestData, timeout=1)
    )
    for _ in range(20):
        if transport.packets:
            break
        await asyncio.sleep(0.01)

    assert len(transport.packets) == 1
    request = transport.packets[0]
    response = ProtoTestData(index=10, message="echo: request")
    communicator.on_received_data(
        type_name_of(ProtoTestData),
        response.SerializeToString(),
        nonce=2,
        response_nonce=request.nonce,
    )

    result = await task
    assert result.index == 10
    assert result.message == "echo: request"

    await communicator.close()


@pytest.mark.asyncio
async def test_nonce_wraps_after_int32_max_without_emitting_zero():
    transport = FakeTransport()
    communicator = Communicator()
    communicator.initialize(transport, parser_map())
    communicator._nonce = MAX_NONCE - 1

    await communicator.send(ProtoTestData(index=1))
    await communicator.send(ProtoTestData(index=2))

    assert [packet.nonce for packet in transport.packets] == [MAX_NONCE, 1]
    assert all(packet.nonce != 0 for packet in transport.packets)

    await communicator.close()


@pytest.mark.asyncio
async def test_send_request_nonce_wraps_at_int32_max():
    transport = FakeTransport()
    communicator = Communicator()
    communicator.initialize(transport, parser_map())
    communicator._nonce = MAX_NONCE - 1

    task = asyncio.create_task(
        communicator.send_request(ProtoTestData(index=1, message="max"), ProtoTestData, timeout=1)
    )
    for _ in range(20):
        if transport.packets:
            break
        await asyncio.sleep(0.01)

    request = transport.packets[0]
    assert request.nonce == MAX_NONCE
    assert request.nonce != 0
    communicator.on_received_data(
        type_name_of(ProtoTestData),
        ProtoTestData(index=2, message="max response").SerializeToString(),
        nonce=1,
        response_nonce=request.nonce,
    )

    result = await task
    assert result.message == "max response"

    await communicator.close()


@pytest.mark.asyncio
async def test_shutdown():
    transport = FakeTransport()
    communicator = Communicator()
    communicator.initialize(transport, parser_map())
    await communicator.close()

    assert not communicator.is_alive()
    assert transport.closed


@pytest.mark.asyncio
async def test_shutdown_public_method():
    transport = FakeTransport()
    communicator = Communicator()
    communicator.initialize(transport, parser_map())
    communicator.shutdown()
    await communicator.wait_closed()
    await asyncio.sleep(0)

    assert not communicator.is_alive()


@pytest.mark.asyncio
async def test_inbound_queue_fifo_order():
    """on_received_data가 여러 메시지를 FIFO 순서대로 전달한다."""
    transport = FakeTransport()
    communicator = Communicator()
    communicator.initialize(transport, parser_map())

    received_indices: list[int] = []
    communicator.add_listener(
        type_name_of(ProtoTestData),
        lambda m: received_indices.append(m.index),
    )

    for i in range(1, 6):
        msg = ProtoTestData(index=i, message=f"msg{i}")
        communicator.on_received_data(type_name_of(ProtoTestData), msg.SerializeToString(), i, 0)

    # receive_loop가 처리할 수 있도록 이벤트 루프를 양보
    for _ in range(20):
        if len(received_indices) == 5:
            break
        await asyncio.sleep(0.01)

    assert received_indices == [1, 2, 3, 4, 5], f"FIFO violated: {received_indices}"
    await communicator.close()


@pytest.mark.asyncio
async def test_slow_request_handler_response_fifo():
    """느린 request handler라도 자동 응답 순서가 FIFO를 보존한다."""
    transport = FakeTransport()
    communicator = Communicator()
    communicator.initialize(transport, parser_map())

    async def slow_handler(req: ProtoTestData, nonce: int) -> ProtoTestData:
        if req.index == 1:
            await asyncio.sleep(0.05)  # 첫 번째 요청 50ms 지연
        return ProtoTestData(index=req.index + 100, message="echo")

    communicator.add_request_listener(type_name_of(ProtoTestData), slow_handler)

    msg1 = ProtoTestData(index=1, message="first")
    msg2 = ProtoTestData(index=2, message="second")
    communicator.on_received_data(type_name_of(ProtoTestData), msg1.SerializeToString(), 10, 0)
    communicator.on_received_data(type_name_of(ProtoTestData), msg2.SerializeToString(), 20, 0)

    for _ in range(40):
        if len(transport.packets) >= 2:
            break
        await asyncio.sleep(0.01)

    assert len(transport.packets) >= 2, f"expected 2 responses, got {len(transport.packets)}"
    # FIFO: 첫 번째 응답은 responseNonce=10, 두 번째는 responseNonce=20
    assert transport.packets[0].responseNonce == 10, f"expected 10, got {transport.packets[0].responseNonce}"
    assert transport.packets[1].responseNonce == 20, f"expected 20, got {transport.packets[1].responseNonce}"
    await communicator.close()


@pytest.mark.asyncio
async def test_close_cancels_pending_receive():
    """close 시 inbound queue에 쌓인 처리 중인 항목이 정리된다."""
    transport = FakeTransport()
    communicator = Communicator()
    communicator.initialize(transport, parser_map())

    # send_request로 pending 상태를 만든다
    task = asyncio.create_task(
        communicator.send_request(ProtoTestData(index=1), ProtoTestData, timeout=5.0)
    )
    await asyncio.sleep(0.01)

    await communicator.close()

    with pytest.raises((Exception,)):
        await task

    assert not communicator.is_alive()


@pytest.mark.asyncio
async def test_inbound_backpressure_on_heavy_load():
    """bounded inbound queue가 포화되었을 때 enqueue_inbound가 블로킹되며 backpressure를 가동하는지 검증한다."""
    transport = FakeTransport()
    communicator = Communicator()
    communicator.initialize(transport, parser_map())
    handler_promise = asyncio.Event()

    async def blocking_handler(msg, nonce=0):
        await handler_promise.wait()
        return ProtoTestData(index=100)

    communicator.add_request_listener("TestData", blocking_handler)

    msg = ProtoTestData(index=1, message="heavy")
    data = msg.SerializeToString()

    # 1번째 호출: 핸들러가 handler_promise에 의해 블로킹됨
    await communicator.enqueue_inbound("TestData", data, 1, 0)

    # 2번째부터 65번째(총 64개)를 enqueue: 첫번째 아이템은 디스패치 중이므로 큐에는 현재 0개 있음
    # 2 ~ 65번째(총 64개) 추가 시 큐가 꽉 차게 됨
    for i in range(2, 66):
        await communicator.enqueue_inbound("TestData", data, i, 0)

    # 66번째 item 추가 시도: 비동기로 대기(블로킹)해야 함
    enqueue_task = asyncio.create_task(communicator.enqueue_inbound("TestData", data, 66, 0))
    await asyncio.sleep(0.02)

    # 큐가 비워지지 않았으므로 enqueue_task는 unblock되지 않아야 함 (Pending)
    assert not enqueue_task.done(), "enqueue_inbound should block when the queue is full"

    # 핸들러를 완료시켜 큐를 하나 소모하게 함
    handler_promise.set()
    await asyncio.sleep(0.01)

    # 대기 중이던 enqueue_task가 무사히 완료되어야 함
    assert enqueue_task.done(), "enqueue_inbound should resume when space becomes available"

    communicator.shutdown()
    await communicator.wait_closed()


@pytest.mark.asyncio
async def test_graceful_close_drains_in_flight():
    """slow request handler가 실행 중인 상태에서 close()를 호출해도 handler가 끝까지 무사히 동작하여 응답을 송신하는지 검증한다."""
    transport = FakeTransport()
    communicator = Communicator()
    communicator.initialize(transport, parser_map())

    handler_started = asyncio.Event()
    handler_finished = asyncio.Event()

    async def slow_handler(req: ProtoTestData) -> ProtoTestData:
        handler_started.set()
        await asyncio.sleep(0.1)
        handler_finished.set()
        return ProtoTestData(index=100, message="graceful")

    communicator.add_request_listener(type_name_of(ProtoTestData), slow_handler)

    msg = ProtoTestData(index=1, message="test")
    communicator.on_received_data(type_name_of(ProtoTestData), msg.SerializeToString(), 1, 0)

    # 핸들러가 구동될 때까지 대기
    await handler_started.wait()

    # 이 상태에서 close() 호출
    close_task = asyncio.create_task(communicator.close())

    # close() 가 완료될 때까지 기다리기 전에 핸들러가 무사히 완주하여 finishes 되었는지 확인
    await handler_finished.wait()
    await close_task

    # 결과 검증: close 완료 후에도 핸들러 응답이 정상적으로 전송 큐에 잘 적재되었는지 확인
    assert len(transport.packets) == 1
    assert transport.packets[0].responseNonce == 1
    assert not communicator.is_alive()

from dataclasses import dataclass

@dataclass
class FakeGatewayItem:
    seq: int
    type_name: str
    data: bytes
    incoming_nonce: int

@pytest.mark.asyncio
async def test_worker_gateway_reorder_and_dispatch():
    """Worker gateway에서 out-of-order로 전송되어도 seq 순으로 정렬 후 coordinator로 밀어넣어 순서를 보존하는지 검증한다."""
    transport = FakeTransport()
    communicator = Communicator()
    communicator.initialize(transport, parser_map())

    received_indices: list[int] = []
    communicator.add_listener(
        type_name_of(ProtoTestData),
        lambda m: received_indices.append(m.index),
    )

    # Fake gateway: out-of-order 아이템들
    items = [
        FakeGatewayItem(
            seq=2,
            type_name=type_name_of(ProtoTestData),
            data=ProtoTestData(index=2).SerializeToString(),
            incoming_nonce=2,
        ),
        FakeGatewayItem(
            seq=1,
            type_name=type_name_of(ProtoTestData),
            data=ProtoTestData(index=1).SerializeToString(),
            incoming_nonce=1,
        )
    ]

    # Reorder by seq
    items.sort(key=lambda x: x.seq)

    # Dispatch to coordinator in reordered sequence
    for item in items:
        communicator.on_received_data(
            item.type_name,
            item.data,
            item.incoming_nonce,
            0,
        )

    # receive_loop가 처리할 수 있도록 이벤트 루프 양보
    for _ in range(20):
        if len(received_indices) == 2:
            break
        await asyncio.sleep(0.01)

    assert received_indices == [1, 2], f"seq ordering violated: {received_indices}"
    await communicator.close()