openttd-client/tests/test_admin.py

import pytest
import asyncio
import json
import os
import monocypher
from openttd import OpenTTDAdminClient
from openttd.protocol import PacketAdminType, AdminUpdateType, AdminUpdateFrequency

def decode_gamescript_payload(packet):
    """Decode the JSON payload of an AdminGamescript packet (2-byte length + 1-byte type + string)."""
    return json.loads(packet[3:].split(b"\x00")[0])

class MockTransport:
    def __init__(self):
        self._closing = False
    def is_closing(self):
        return self._closing
    def close(self):
        self._closing = True
    def write(self, data):
        return len(data)

class MockProtocol:
    def __init__(self):
        self.sent = []
    async def send_packet(self, data):
        self.sent.append(data)
        return len(data)

def test_admin_packet_types():
    assert PacketAdminType.AdminJoin == 0
    assert PacketAdminType.ServerWelcome == 104
    assert PacketAdminType.ServerAuthRequest == 128

@pytest.mark.asyncio
async def test_admin_list_vehicles():
    client = OpenTTDAdminClient("127.0.0.1", port=3977, admin_name="TestAdmin")
    proto = MockProtocol()
    client._protocol = proto
    client._transport = MockTransport()

    await client.list_vehicles()
    await client.list_vehicles(company_id=2)

    assert len(proto.sent) == 2
    assert decode_gamescript_payload(proto.sent[0]) == {"command": "list_vehicles"}
    assert decode_gamescript_payload(proto.sent[1]) == {"command": "list_vehicles", "company_id": 2}

@pytest.mark.asyncio
async def test_admin_client_connect_and_actions(monkeypatch):
    client = OpenTTDAdminClient("127.0.0.1", port=3977, admin_name="TestAdmin")
    assert client.host == "127.0.0.1"
    assert client.port == 3977
    assert client.admin_name == "TestAdmin"
    
    # 1. Connect (secure=False)
    proto = MockProtocol()
    transport = MockTransport()
    async def mock_connect(*args, **kwargs):
        return transport, proto
    
    monkeypatch.setattr(asyncio.get_running_loop(), "create_connection", mock_connect)
    await client.connect(admin_password="asd", secure=False)
    assert len(proto.sent) == 1
    
    # 2. Connect (secure=True)
    proto.sent.clear()
    await client.connect(admin_password="asd", secure=True)
    assert len(proto.sent) == 1

    # 3. Connect Exception
    async def mock_fail(*args, **kwargs):
        raise Exception("Connection Failed")
    monkeypatch.setattr(asyncio.get_running_loop(), "create_connection", mock_fail)
    with pytest.raises(Exception, match="Connection Failed"):
        await client.connect()

    # Reset connection mock for further tests
    client._protocol = proto
    client._transport = transport
    
    # 4. Actions
    proto.sent.clear()
    await client.send_rcon("help")
    await client.send_chat("hello")
    await client.update_frequency(AdminUpdateType.Chat, AdminUpdateFrequency.Automatic)
    await client.poll(AdminUpdateType.ClientInfo)
    await client.poll_clients(1)
    await client.poll_companies(2)
    await client.poll_economy(3)
    await client.poll_stats(4)
    await client.send_gamescript({"cmd": "test"})
    assert len(proto.sent) == 9

    # 5. Protocol callbacks
    client.connected(None)
    
    # X25519 authentication handshake callback
    server_pub = monocypher.x25519_public_key(monocypher.generate_key())
    nonce = os.urandom(24)
    await client.receive_ServerAuthRequest(None, 1, server_pub + nonce)
    assert client._session_key_send is not None
    
    await client.receive_ServerEnableEncryption(None, os.urandom(24))
    assert client.encryption_enabled is True
    
    await client.receive_ServerProtocol(None, 3, {})
    await client.receive_ServerWelcome(None, server_name="JGRPP")
    assert client.joined.is_set()

    # Chat callback with and without handler
    chat_events = []
    client.on_chat = lambda **k: chat_events.append(k)
    await client.receive_ServerChat(None, client_id=1, message="hello")
    assert chat_events == [{"client_id": 1, "message": "hello"}]
    client.on_chat = None
    await client.receive_ServerChat(None, client_id=1, message="hello")

    # Console callback with and without handler
    console_events = []
    client.on_console = lambda **k: console_events.append(k)
    await client.receive_ServerConsole(None, origin="server", text="welcome")
    assert console_events == [{"origin": "server", "text": "welcome"}]
    client.on_console = None
    await client.receive_ServerConsole(None, origin="server", text="welcome")

    # Other callbacks
    await client.receive_ServerRcon(None, text="res")
    await client.receive_ServerRconEnd(None, command="help")
    await client.receive_ServerClientJoin(None, client_id=1)
    await client.receive_ServerClientInfo(None, client_id=1, name="user", network_address="127.0.0.1")
    await client.receive_ServerClientUpdate(None, client_id=1, name="user2")
    await client.receive_ServerClientQuit(None, client_id=1)
    await client.receive_ServerClientError(None, client_id=1, error_code=3)
    await client.receive_ServerCompanyNew(None, company_id=1)
    await client.receive_ServerCompanyInfo(None, company_id=1, name="company")
    await client.receive_ServerCompanyUpdate(None, company_id=1)
    await client.receive_ServerCompanyRemove(None, company_id=1)
    await client.receive_ServerCompanyEconomy(None, company_id=1, money=1000, loan=100)
    await client.receive_ServerCompanyStats(None, company_id=1, vehicles={}, stations={})
    
    # Gamescript callback with and without handler
    gs_events = []
    client.on_gamescript = lambda data: gs_events.append(data)
    await client.receive_ServerGamescript(None, data={"a": 1})
    assert gs_events == [{"a": 1}]
    client.on_gamescript = None
    await client.receive_ServerGamescript(None, data={"a": 1})

    await client.receive_ServerDate(None)
    
    # Error callbacks that shut down the client
    client.shutdown_event.clear()
    await client.receive_ServerError(None, 5)
    assert client.shutdown_event.is_set()
    
    client.shutdown_event.clear()
    await client.receive_ServerFull(None)
    assert client.shutdown_event.is_set()
    
    client.shutdown_event.clear()
    await client.receive_ServerBanned(None)
    assert client.shutdown_event.is_set()
    
    client.shutdown_event.clear()
    await client.receive_ServerShutdown(None)
    assert client.shutdown_event.is_set()
    
    await client.receive_ServerNewGame(None)
    await client.receive_ServerPong(None, payload=123)

    # 6. Disconnect
    client.shutdown_event.clear()
    client.disconnect(None)
    assert client.shutdown_event.is_set()

    # 7. Quit
    client._transport = MockTransport()
    await client.quit()
    assert client.shutdown_event.is_set()

    # 8. Quit Exception
    class BadProtocol:
        async def send_packet(self, data):
            raise Exception("Fail")
    client._transport = MockTransport()
    client._protocol = BadProtocol()
    await client.quit()
    assert client.shutdown_event.is_set()