113 lines
3.1 KiB
Python
113 lines
3.1 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
|
|
import pytest
|
|
|
|
from toki_socket.communicator import 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_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()
|