290 lines
9.6 KiB
Python
290 lines
9.6 KiB
Python
import asyncio
|
|
import unittest
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
from arkteos_proxy import (
|
|
ArkteosProxy,
|
|
ConfigurationError,
|
|
KEEPALIVE_INTERVAL,
|
|
RECONNECT_DELAY,
|
|
validate_configuration,
|
|
)
|
|
|
|
|
|
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
|
|
proxy.wait_for_reconnect_delay = AsyncMock()
|
|
with (
|
|
patch("arkteos_proxy.asyncio.start_server", new_callable=AsyncMock, return_value=FakeServer()),
|
|
):
|
|
await proxy.serve()
|
|
|
|
self.assertEqual(attempts, 2)
|
|
proxy.wait_for_reconnect_delay.assert_awaited_once_with()
|
|
|
|
def test_missing_option_defaults_to_read_only(self):
|
|
self.assertFalse(ArkteosProxy("pac", 9641, 9641).allow_client_writes)
|
|
|
|
async def test_connection_timeout_stops_a_single_attempt(self):
|
|
proxy = ArkteosProxy("pac", 9641, 9641)
|
|
|
|
async def never_connect(*_args, **_kwargs):
|
|
await asyncio.Event().wait()
|
|
|
|
with (
|
|
patch("arkteos_proxy.PAC_CONNECT_TIMEOUT", 0.01),
|
|
patch("arkteos_proxy.asyncio.open_connection", new_callable=AsyncMock, side_effect=never_connect) as connect,
|
|
):
|
|
with self.assertRaises(ConnectionError):
|
|
await proxy.run_pac_connection()
|
|
|
|
connect.assert_awaited_once_with("pac", 9641)
|
|
|
|
async def test_silent_pac_closes_writer_and_cancels_keepalive(self):
|
|
proxy = ArkteosProxy("pac", 9641, 9641)
|
|
reader = QueueReader()
|
|
writer = FakeWriter()
|
|
keepalive_cancelled = asyncio.Event()
|
|
|
|
async def tracked_keepalive():
|
|
try:
|
|
await asyncio.Event().wait()
|
|
except asyncio.CancelledError:
|
|
keepalive_cancelled.set()
|
|
raise
|
|
|
|
proxy.pac_keepalive = tracked_keepalive
|
|
with (
|
|
patch("arkteos_proxy.PAC_READ_TIMEOUT", 0.01),
|
|
patch("arkteos_proxy.asyncio.open_connection", new_callable=AsyncMock, return_value=(reader, writer)),
|
|
):
|
|
await proxy.run_pac_connection()
|
|
|
|
self.assertTrue(writer.closed)
|
|
self.assertTrue(keepalive_cancelled.is_set())
|
|
self.assertIsNone(proxy.pac_writer)
|
|
self.assertIsNone(proxy.keepalive_task)
|
|
|
|
async def test_pac_activity_is_distributed_before_connection_closes(self):
|
|
proxy = ArkteosProxy("pac", 9641, 9641)
|
|
reader = QueueReader()
|
|
writer = FakeWriter()
|
|
client = FakeWriter()
|
|
proxy.clients.add(client)
|
|
await reader.items.put(b"pac-data")
|
|
await reader.items.put(b"")
|
|
|
|
with (
|
|
patch("arkteos_proxy.PAC_READ_TIMEOUT", 0.01),
|
|
patch("arkteos_proxy.asyncio.open_connection", new_callable=AsyncMock, return_value=(reader, writer)),
|
|
):
|
|
await proxy.run_pac_connection()
|
|
|
|
self.assertEqual(client.writes, [b"pac-data"])
|
|
self.assertTrue(writer.closed)
|
|
|
|
async def test_graceful_stop_closes_server_clients_and_pac_writer(self):
|
|
proxy = ArkteosProxy("pac", 9641, 0)
|
|
server = FakeServer()
|
|
pac_writer = FakeWriter()
|
|
client = FakeWriter()
|
|
proxy.pac_writer = pac_writer
|
|
proxy.clients.add(client)
|
|
|
|
async def wait_for_stop():
|
|
await proxy.stop_event.wait()
|
|
|
|
proxy.run_pac_connection = wait_for_stop
|
|
with patch("arkteos_proxy.asyncio.start_server", new_callable=AsyncMock, return_value=server):
|
|
task = asyncio.create_task(proxy.serve())
|
|
await asyncio.sleep(0)
|
|
proxy.request_stop()
|
|
await task
|
|
|
|
self.assertTrue(server.closed)
|
|
self.assertTrue(client.closed)
|
|
self.assertTrue(pac_writer.closed)
|
|
self.assertEqual(proxy.client_tasks, set())
|
|
|
|
async def test_stop_during_reconnect_does_not_wait_for_full_delay(self):
|
|
proxy = ArkteosProxy("pac", 9641, 0)
|
|
server = FakeServer()
|
|
|
|
async def failed_connection():
|
|
raise ConnectionError("PAC indisponible")
|
|
|
|
async def stop_during_delay():
|
|
proxy.request_stop()
|
|
|
|
proxy.run_pac_connection = failed_connection
|
|
proxy.wait_for_reconnect_delay = stop_during_delay
|
|
with patch("arkteos_proxy.asyncio.start_server", new_callable=AsyncMock, return_value=server):
|
|
await asyncio.wait_for(proxy.serve(), timeout=0.1)
|
|
|
|
self.assertTrue(server.closed)
|
|
|
|
def test_configuration_validation_rejects_invalid_values(self):
|
|
invalid_configurations = (
|
|
(" ", "9641", "9641", "false"),
|
|
("pac", "0", "9641", "false"),
|
|
("pac", "65536", "9641", "false"),
|
|
("pac", "9641", "invalid", "false"),
|
|
("pac", "9641", "9641", "maybe"),
|
|
)
|
|
for configuration in invalid_configurations:
|
|
with self.subTest(configuration=configuration):
|
|
with self.assertRaises(ConfigurationError):
|
|
validate_configuration(*configuration)
|
|
|
|
def test_configuration_validation_accepts_supported_values(self):
|
|
self.assertEqual(
|
|
validate_configuration(" pac ", "9641", 9641, "false"),
|
|
("pac", 9641, 9641, False),
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|