openttd-client/tests/test_admin.py

import asyncio
import json
import os

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


class FakeNetworkError(OSError):
    """Stand-in for a socket-level failure, so it matches the client's narrowed handlers."""

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_list_stations():
    client = OpenTTDAdminClient("127.0.0.1", port=3977, admin_name="TestAdmin")
    proto = MockProtocol()
    client._protocol = proto
    client._transport = MockTransport()

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

    assert len(proto.sent) == 2
    assert decode_gamescript_payload(proto.sent[0]) == {"command": "list_stations"}
    assert decode_gamescript_payload(proto.sent[1]) == {"command": "list_stations", "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 FakeNetworkError("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 FakeNetworkError("Fail")
    client._transport = MockTransport()
    client._protocol = BadProtocol()
    await client.quit()
    assert client.shutdown_event.is_set()

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

    task = asyncio.ensure_future(client.get_timetable(5))
    await asyncio.sleep(0)  # let the task send the request

    # First use auto-subscribes to Gamescript updates, then sends the query.
    assert len(proto.sent) == 2
    assert proto.sent[0][2] == PacketAdminType.AdminUpdateFrequency
    assert decode_gamescript_payload(proto.sent[1]) == {
        "command": "get_timetable", "vehicle_id": 5, "request_id": 1,
    }

    response = {"command": "get_timetable", "vehicle_id": 5, "request_id": 1,
                "lateness": 0, "start_tick": 0, "current_order_time": 3,
                "total_duration": -1, "orders": []}
    await client.receive_ServerGamescript(None, data=response)
    assert await task == response
    assert client._gs_futures == {}

    # Second call must not re-subscribe and must use a fresh request id.
    task = asyncio.ensure_future(client.get_timetable(9))
    await asyncio.sleep(0)
    assert len(proto.sent) == 3
    assert decode_gamescript_payload(proto.sent[2])["request_id"] == 2
    await client.receive_ServerGamescript(None, data={"request_id": 2, "orders": []})
    assert (await task)["orders"] == []

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

    with pytest.raises(asyncio.TimeoutError):
        await client.get_timetable(5, timeout=0.05)
    assert client._gs_futures == {}

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

    task = asyncio.ensure_future(client.get_timetable(65535))
    await asyncio.sleep(0)
    await client.receive_ServerGamescript(
        None, data={"command": "get_timetable", "vehicle_id": 65535,
                    "request_id": 1, "error": "invalid_vehicle"})
    with pytest.raises(ValueError, match="invalid_vehicle"):
        await task
    assert client._gs_futures == {}

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

    task = asyncio.ensure_future(client.get_timetable(5))
    await asyncio.sleep(0)
    client.disconnect(None)
    with pytest.raises(ConnectionError):
        await task
    assert client._gs_futures == {}

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

    task = asyncio.ensure_future(client.get_station(3))
    await asyncio.sleep(0)  # let the task send the request

    # First use auto-subscribes to Gamescript updates, then sends the query.
    assert len(proto.sent) == 2
    assert proto.sent[0][2] == PacketAdminType.AdminUpdateFrequency
    assert decode_gamescript_payload(proto.sent[1]) == {
        "command": "get_station", "station_id": 3, "request_id": 1,
    }

    response = {"command": "get_station", "station_id": 3, "request_id": 1,
                "name": "Test Central", "location": 12345, "owner": 0,
                "cargo": [{"cargo_id": 0, "waiting": 42, "planned": 17, "rating": 71}]}
    await client.receive_ServerGamescript(None, data=response)
    assert await task == response
    assert client._gs_futures == {}

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

    task = asyncio.ensure_future(client.get_station(65535))
    await asyncio.sleep(0)
    await client.receive_ServerGamescript(
        None, data={"command": "get_station", "station_id": 65535,
                    "request_id": 1, "error": "invalid_station"})
    with pytest.raises(ValueError, match="invalid_station"):
        await task
    assert client._gs_futures == {}

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

    with pytest.raises(asyncio.TimeoutError):
        await client.get_station(3, timeout=0.05)
    assert client._gs_futures == {}

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

    task = asyncio.ensure_future(client.get_station_cargo(3, 0))
    await asyncio.sleep(0)  # let the task send the request

    # First use auto-subscribes to Gamescript updates, then sends the query.
    assert len(proto.sent) == 2
    assert proto.sent[0][2] == PacketAdminType.AdminUpdateFrequency
    assert decode_gamescript_payload(proto.sent[1]) == {
        "command": "get_station_cargo", "station_id": 3, "cargo_id": 0, "request_id": 1,
    }

    response = {"command": "get_station_cargo", "station_id": 3, "cargo_id": 0, "request_id": 1,
                "waiting": 60, "planned": 40,
                "waiting_by_from": [{"station": 5, "amount": 25}, {"station": 6, "amount": 35}],
                "planned_by_from": [{"station": 5, "amount": 40}],
                "waiting_by_via": [{"station": 7, "amount": 60}],
                "planned_by_via": [{"station": 7, "amount": 40}]}
    await client.receive_ServerGamescript(None, data=response)
    assert await task == response
    assert client._gs_futures == {}

@pytest.mark.asyncio
async def test_admin_get_station_cargo_with_filters_encoding():
    client = OpenTTDAdminClient("127.0.0.1", port=3977, admin_name="TestAdmin")
    proto = MockProtocol()
    client._protocol = proto
    client._transport = MockTransport()
    client._gs_subscribed = True  # skip the auto-subscribe so only the query is sent

    task = asyncio.ensure_future(client.get_station_cargo(3, 0, from_station=6, via_station=7))
    await asyncio.sleep(0)

    assert len(proto.sent) == 1
    assert decode_gamescript_payload(proto.sent[0]) == {
        "command": "get_station_cargo", "station_id": 3, "cargo_id": 0,
        "from_station": 6, "via_station": 7, "request_id": 1,
    }
    await client.receive_ServerGamescript(
        None, data={"request_id": 1, "waiting": 12, "planned": 8})
    assert (await task)["waiting"] == 12

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

    task = asyncio.ensure_future(client.get_station_cargo(3, 999))
    await asyncio.sleep(0)
    await client.receive_ServerGamescript(
        None, data={"command": "get_station_cargo", "station_id": 3, "cargo_id": 999,
                    "request_id": 1, "error": "invalid_cargo"})
    with pytest.raises(ValueError, match="invalid_cargo"):
        await task
    assert client._gs_futures == {}

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

    task = asyncio.ensure_future(client.get_dispatch(7))
    await asyncio.sleep(0)  # let the task send the request

    # First use auto-subscribes to Gamescript updates, then sends the query.
    assert len(proto.sent) == 2
    assert proto.sent[0][2] == PacketAdminType.AdminUpdateFrequency
    assert decode_gamescript_payload(proto.sent[1]) == {
        "command": "get_dispatch", "vehicle_id": 7, "request_id": 1,
    }

    response = {"command": "get_dispatch", "vehicle_id": 7, "request_id": 1,
                "enabled": 1,
                "schedules": [{"index": 0, "duration": 3000, "start_tick": 0, "delay": 0,
                               "reuse_slots": 0,
                               "slots": [{"offset": 500, "flags": 0}, {"offset": 1500, "flags": 0}]}]}
    await client.receive_ServerGamescript(None, data=response)
    assert await task == response
    assert client._gs_futures == {}

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

    task = asyncio.ensure_future(client.get_dispatch(65535))
    await asyncio.sleep(0)
    await client.receive_ServerGamescript(
        None, data={"command": "get_dispatch", "vehicle_id": 65535,
                    "request_id": 1, "error": "invalid_vehicle"})
    with pytest.raises(ValueError, match="invalid_vehicle"):
        await task
    assert client._gs_futures == {}

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

    gs_events = []
    client.on_gamescript = lambda data: gs_events.append(data)

    # No request_id, unknown request_id, and non-dict payloads all pass through.
    await client.receive_ServerGamescript(None, data={"vehicles": []})
    await client.receive_ServerGamescript(None, data={"request_id": 999, "orders": []})
    await client.receive_ServerGamescript(None, data="plain string")
    assert gs_events == [{"vehicles": []}, {"request_id": 999, "orders": []}, "plain string"]