Extends the AdminBridge GameScript JSON channel (the same relay used by list_vehicles/get_timetable) with station support: - list_stations(): enumerate stations, fire-and-forget like list_vehicles(). - get_station(): authoritative per-cargo snapshot of a station's live state, with both the real-time waiting amount (GSStation.GetCargoWaiting) and the planned cargodist link-graph flow (GetCargoPlanned), plus rating. - get_station_cargo(): break one cargo type down by source station and by next hop (the cargodist routing destination) for both waiting and planned amounts, with optional from_station/via_station filters. All three use stock GameScript API (no server patch, unlike timetables). Refactors the shared GS request/reply correlation out of get_timetable and get_station into a _gs_query() helper. Companion handlers must be added to the server-side AdminBridge GameScript (not tracked in this repo). Includes unit + e2e tests, a worked demo in main_admin.py, and protocol docs. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
400 lines
16 KiB
Python
400 lines
16 KiB
Python
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"]
|