167 lines
5.2 KiB
Python
167 lines
5.2 KiB
Python
import asyncio
|
|
import unittest
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
from arkteos_proxy import ArkteosProxy, KEEPALIVE_INTERVAL, RECONNECT_DELAY
|
|
|
|
|
|
class QueueReader:
|
|
def __init__(self):
|
|
self.items = asyncio.Queue()
|
|
|
|
async def read(self, _size):
|
|
return await self.items.get()
|
|
|
|
|
|
class FakeWriter:
|
|
def __init__(self):
|
|
self.writes = []
|
|
self.drain_calls = 0
|
|
self.closed = False
|
|
|
|
def write(self, data):
|
|
self.writes.append(data)
|
|
|
|
async def drain(self):
|
|
self.drain_calls += 1
|
|
|
|
def close(self):
|
|
self.closed = True
|
|
|
|
async def wait_closed(self):
|
|
return None
|
|
|
|
def get_extra_info(self, _name):
|
|
return ("127.0.0.1", 12345)
|
|
|
|
|
|
class SerialWriter(FakeWriter):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.active_drains = 0
|
|
self.maximum_active_drains = 0
|
|
|
|
async def drain(self):
|
|
self.drain_calls += 1
|
|
self.active_drains += 1
|
|
self.maximum_active_drains = max(self.maximum_active_drains, self.active_drains)
|
|
await asyncio.sleep(0)
|
|
self.active_drains -= 1
|
|
|
|
|
|
class FakeServer:
|
|
def __init__(self):
|
|
self.closed = False
|
|
|
|
def close(self):
|
|
self.closed = True
|
|
|
|
async def wait_closed(self):
|
|
return None
|
|
|
|
|
|
class ArkteosProxyTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_default_mode_blocks_client_writes_and_keeps_client_connected(self):
|
|
proxy = ArkteosProxy("pac", 9641, 9641)
|
|
pac_writer = FakeWriter()
|
|
proxy.pac_writer = pac_writer
|
|
reader = QueueReader()
|
|
client_writer = FakeWriter()
|
|
task = asyncio.create_task(proxy.handle_client(reader, client_writer))
|
|
|
|
await reader.items.put(b"client-data")
|
|
await asyncio.sleep(0)
|
|
self.assertEqual(pac_writer.writes, [])
|
|
self.assertIn(client_writer, proxy.clients)
|
|
|
|
await proxy.broadcast_to_clients(b"pac-data")
|
|
self.assertEqual(client_writer.writes, [b"pac-data"])
|
|
|
|
await reader.items.put(b"")
|
|
await task
|
|
|
|
async def test_bidirectional_mode_relays_data_and_drains_writer(self):
|
|
proxy = ArkteosProxy("pac", 9641, 9641, allow_client_writes=True)
|
|
pac_writer = FakeWriter()
|
|
proxy.pac_writer = pac_writer
|
|
reader = QueueReader()
|
|
client_writer = FakeWriter()
|
|
task = asyncio.create_task(proxy.handle_client(reader, client_writer))
|
|
|
|
await reader.items.put(b"client-data")
|
|
await reader.items.put(b"")
|
|
await task
|
|
|
|
self.assertEqual(pac_writer.writes, [b"client-data"])
|
|
self.assertEqual(pac_writer.drain_calls, 1)
|
|
|
|
async def test_two_clients_writes_are_serialized(self):
|
|
proxy = ArkteosProxy("pac", 9641, 9641, allow_client_writes=True)
|
|
pac_writer = SerialWriter()
|
|
proxy.pac_writer = pac_writer
|
|
first_reader = QueueReader()
|
|
second_reader = QueueReader()
|
|
first_task = asyncio.create_task(proxy.handle_client(first_reader, FakeWriter()))
|
|
second_task = asyncio.create_task(proxy.handle_client(second_reader, FakeWriter()))
|
|
|
|
await first_reader.items.put(b"first")
|
|
await second_reader.items.put(b"second")
|
|
await first_reader.items.put(b"")
|
|
await second_reader.items.put(b"")
|
|
await asyncio.gather(first_task, second_task)
|
|
|
|
self.assertCountEqual(pac_writer.writes, [b"first", b"second"])
|
|
self.assertEqual(pac_writer.maximum_active_drains, 1)
|
|
|
|
async def test_keepalive_uses_serialized_write_and_keeps_protocol_values(self):
|
|
proxy = ArkteosProxy("pac", 9641, 9641)
|
|
pac_writer = SerialWriter()
|
|
proxy.pac_writer = pac_writer
|
|
|
|
await proxy.send_keepalive()
|
|
|
|
self.assertEqual(KEEPALIVE_INTERVAL, 300)
|
|
self.assertEqual(pac_writer.writes, [b"\x00"])
|
|
self.assertEqual(pac_writer.drain_calls, 1)
|
|
self.assertEqual(pac_writer.maximum_active_drains, 1)
|
|
|
|
async def test_pac_disconnect_closes_clients(self):
|
|
proxy = ArkteosProxy("pac", 9641, 9641)
|
|
first_client = FakeWriter()
|
|
second_client = FakeWriter()
|
|
proxy.clients.update({first_client, second_client})
|
|
|
|
await proxy.close_clients()
|
|
|
|
self.assertTrue(first_client.closed)
|
|
self.assertTrue(second_client.closed)
|
|
self.assertEqual(proxy.clients, set())
|
|
|
|
async def test_reconnection_loop_is_preserved(self):
|
|
proxy = ArkteosProxy("pac", 9641, 0)
|
|
attempts = 0
|
|
|
|
async def fake_pac_connection():
|
|
nonlocal attempts
|
|
attempts += 1
|
|
if attempts == 1:
|
|
raise ConnectionError("PAC fermée")
|
|
proxy.stop_event.set()
|
|
|
|
proxy.run_pac_connection = fake_pac_connection
|
|
with (
|
|
patch("arkteos_proxy.asyncio.start_server", new_callable=AsyncMock, return_value=FakeServer()),
|
|
patch("arkteos_proxy.asyncio.sleep", new_callable=AsyncMock) as sleep,
|
|
):
|
|
await proxy.serve()
|
|
|
|
self.assertEqual(attempts, 2)
|
|
sleep.assert_awaited_once_with(RECONNECT_DELAY)
|
|
|
|
def test_missing_option_defaults_to_read_only(self):
|
|
self.assertFalse(ArkteosProxy("pac", 9641, 9641).allow_client_writes)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|