proto-socket/python/test/test_communicator.py
toki 9cc1f1d58f sync: update communicator implementation across all languages
- Align protocol documentation (PROTOCOL.md, README.md, VERSIONING.md)
- Go: add nonce test, update communicator
- Kotlin: update Communicator, TcpClient, TcpServer, add TLS test
- Python: update all modules, add certificate test resources
- TypeScript: update communicator, tcp/ws clients and servers, add tests
- Dart: update communicator, heartbeat mixin, and tests
2026-04-26 05:31:56 +09:00

160 lines
4.5 KiB
Python

from __future__ import annotations
import asyncio
import pytest
from toki_socket.communicator import MAX_NONCE, Communicator, type_name_of
from toki_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()