arkteos-proxy-addon/tests/test_arkteos_proxy.py

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()