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_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 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() @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_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"]