openttd-client/tests/test_timetable.py

import pytest
from openttd import OpenTTDClient
from openttd.protocol import (
    INVALID_VEH_ORDER_ID,
    GameCommand,
    ModifyTimetableCtrlFlag,
    ModifyTimetableFlags,
    OpenTTDProtocol,
    OrderNonStopFlags,
    OrderStopLocation,
    OrderType,
    read_varuint,
    read_varuint_signed,
    write_varuint,
    write_varuint_signed,
)
from openttd_protocol.wire.read import read_uint8, read_uint16


class MockTransport:
    def is_closing(self):
        return False
    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 decode_sent_command(packet):
    """Decode a ClientCommand packet by reusing the ServerCommand parser, padding on the
    frame/my_cmd trailer that only ServerCommand carries on the wire (ClientCommand doesn't)."""
    padded = bytes(packet)[3:] + b"\x00\x00\x00\x00\x00"
    return OpenTTDProtocol.receive_ServerCommand(None, memoryview(padded))

def new_client():
    client = OpenTTDClient("127.0.0.1")
    client._protocol = MockProtocol()
    client._transport = MockTransport()
    client._target_company = 0
    return client


# --- Varuint codec ---

@pytest.mark.parametrize("value", [
    0, 1, 127,
    128, 16383,
    16384, 2097151,
    2097152, 268435455,
    268435456, 34359738367,
    34359738368, 4398046511103,
    4398046511104, 562949953421311,
    562949953421312, 72057594037927935,
    72057594037927936, 18446744073709551615,
])
def test_varuint_roundtrip_boundaries(value):
    buf = bytearray()
    write_varuint(buf, value)
    decoded, rest = read_varuint(memoryview(bytes(buf)))
    assert decoded == value
    assert bytes(rest) == b""

def test_varuint_rejects_negative():
    with pytest.raises(ValueError):
        write_varuint(bytearray(), -1)

@pytest.mark.parametrize("value", [0, 1, -1, 2, -2, 1000000, -1000000, 9223372036854775807, -9223372036854775808])
def test_varuint_signed_roundtrip(value):
    buf = bytearray()
    write_varuint_signed(buf, value)
    decoded, rest = read_varuint_signed(memoryview(bytes(buf)))
    assert decoded == value
    assert bytes(rest) == b""


# --- OpenTTDProtocol.receive_ServerCommand ---

def build_server_command_bytes(company, cmd, payload, callback=0, callback_param=0, frame=42, my_cmd=True):
    import struct
    body = bytearray()
    body += struct.pack("<B", company)
    body += struct.pack("<H", cmd)
    body += struct.pack("<H", 0)
    body += struct.pack("<I", 0)
    body += struct.pack("<H", len(payload))
    body += payload
    body += struct.pack("<B", callback)
    if callback != 0:
        body += struct.pack("<I", callback_param)
    body += struct.pack("<I", frame)
    body += struct.pack("<B", 1 if my_cmd else 0)
    return bytes(body)

def test_protocol_receive_server_command_no_callback():
    payload = bytearray()
    write_varuint(payload, 7)
    data = build_server_command_bytes(1, GameCommand.ChangeTimetable, payload, callback=0, frame=42, my_cmd=True)
    res = OpenTTDProtocol.receive_ServerCommand(None, memoryview(data))
    assert res["company"] == 1
    assert res["cmd"] == GameCommand.ChangeTimetable
    assert res["callback"] == 0
    assert res["callback_param"] == 0
    assert res["frame"] == 42
    assert res["my_cmd"] is True
    assert bytes(res["payload"]) == bytes(payload)

def test_protocol_receive_server_command_with_callback():
    payload = bytearray()
    write_varuint(payload, 9)
    data = build_server_command_bytes(2, GameCommand.SetVehicleOnTime, payload, callback=5, callback_param=999, frame=100, my_cmd=False)
    res = OpenTTDProtocol.receive_ServerCommand(None, memoryview(data))
    assert res["callback"] == 5
    assert res["callback_param"] == 999
    assert res["my_cmd"] is False


# --- OpenTTDClient outgoing command methods ---

@pytest.mark.asyncio
async def test_client_send_command_with_callback_includes_callback_param():
    client = new_client()
    await client._send_command(GameCommand.ChangeTimetable, bytearray(), callback=5)
    parsed = decode_sent_command(client._protocol.sent[0])
    assert parsed["callback"] == 5
    assert parsed["callback_param"] == 0

@pytest.mark.asyncio
async def test_client_change_timetable_sends_expected_payload():
    client = new_client()
    await client.change_timetable(7, 3, ModifyTimetableFlags.WaitTime, 120)
    assert len(client._protocol.sent) == 1
    parsed = decode_sent_command(client._protocol.sent[0])
    assert parsed["cmd"] == GameCommand.ChangeTimetable
    assert parsed["company"] == 0
    vehicle_id, rest = read_varuint(parsed["payload"])
    order_position, rest = read_uint16(rest)
    flag, rest = read_uint8(rest)
    value, rest = read_varuint(rest)
    ctrl_flags, _ = read_uint8(rest)
    assert (vehicle_id, order_position, flag, value, ctrl_flags) == (7, 3, ModifyTimetableFlags.WaitTime, 120, 0)

@pytest.mark.asyncio
async def test_client_change_timetable_clear_field_sets_ctrl_flag():
    client = new_client()
    await client.change_timetable(7, 3, ModifyTimetableFlags.TravelTime, 0, clear_field=True)
    parsed = decode_sent_command(client._protocol.sent[0])
    _, rest = read_varuint(parsed["payload"])
    _, rest = read_uint16(rest)
    flag, rest = read_uint8(rest)
    _, rest = read_varuint(rest)
    ctrl_flags, _ = read_uint8(rest)
    assert flag == ModifyTimetableFlags.TravelTime
    assert ctrl_flags == ModifyTimetableCtrlFlag.ClearField

@pytest.mark.asyncio
async def test_client_autofill_timetable_sends_expected_payload():
    client = new_client()
    await client.autofill_timetable(7, autofill=True, preserve_wait_time=False)
    parsed = decode_sent_command(client._protocol.sent[0])
    assert parsed["cmd"] == GameCommand.AutofillTimetable
    vehicle_id, rest = read_varuint(parsed["payload"])
    autofill, rest = read_uint8(rest)
    preserve_wait_time, _ = read_uint8(rest)
    assert (vehicle_id, autofill, preserve_wait_time) == (7, 1, 0)

@pytest.mark.asyncio
async def test_client_set_timetable_start_sends_expected_payload():
    client = new_client()
    await client.set_timetable_start(7, True, -12345)
    parsed = decode_sent_command(client._protocol.sent[0])
    assert parsed["cmd"] == GameCommand.SetTimetableStart
    vehicle_id, rest = read_varuint(parsed["payload"])
    timetable_all, rest = read_uint8(rest)
    start_date, _ = read_varuint_signed(rest)
    assert (vehicle_id, timetable_all, start_date) == (7, 1, -12345)

@pytest.mark.asyncio
async def test_client_set_vehicle_on_time_sends_expected_payload():
    client = new_client()
    await client.set_vehicle_on_time(7, apply_to_group=True)
    parsed = decode_sent_command(client._protocol.sent[0])
    assert parsed["cmd"] == GameCommand.SetVehicleOnTime
    vehicle_id, rest = read_varuint(parsed["payload"])
    apply_to_group, _ = read_uint8(rest)
    assert (vehicle_id, apply_to_group) == (7, 1)


# --- OpenTTDClient order add/remove commands ---

def _decode_insert_order_payload(payload):
    vehicle_id, rest = read_varuint(payload)
    sel_ord, rest = read_uint16(rest)
    order_type, rest = read_uint8(rest)
    order_flags, rest = read_uint16(rest)
    station, _ = read_uint16(rest)
    return vehicle_id, sel_ord, order_type, order_flags, station

@pytest.mark.asyncio
async def test_client_add_order_appends_by_default():
    client = new_client()
    await client.add_order(7, 6)
    assert len(client._protocol.sent) == 1
    parsed = decode_sent_command(client._protocol.sent[0])
    assert parsed["cmd"] == GameCommand.InsertOrder
    assert parsed["company"] == 0
    vehicle_id, sel_ord, order_type, order_flags, station = _decode_insert_order_payload(parsed["payload"])
    assert vehicle_id == 7
    assert sel_ord == INVALID_VEH_ORDER_ID  # append to the end
    # OT_GOTO_STATION in bits 0-3, far-end stop location in bits 4-5, stop-everywhere non-stop in 6-7
    assert order_type == (OrderType.GotoStation | (OrderStopLocation.PlatformFarEnd << 4))
    assert order_flags == 0
    assert station == 6

@pytest.mark.asyncio
async def test_client_add_order_insert_position_and_nonstop():
    client = new_client()
    await client.add_order(7, 6, before_position=1, non_stop=OrderNonStopFlags.NoStopAtIntermediate)
    parsed = decode_sent_command(client._protocol.sent[0])
    _, sel_ord, order_type, _, _ = _decode_insert_order_payload(parsed["payload"])
    assert sel_ord == 1  # insert before order position 1
    assert order_type == (OrderType.GotoStation
                          | (OrderStopLocation.PlatformFarEnd << 4)
                          | (OrderNonStopFlags.NoStopAtIntermediate << 6))

@pytest.mark.asyncio
async def test_client_remove_order_sends_expected_payload():
    client = new_client()
    await client.remove_order(7, 2)
    parsed = decode_sent_command(client._protocol.sent[0])
    assert parsed["cmd"] == GameCommand.DeleteOrder
    vehicle_id, rest = read_varuint(parsed["payload"])
    order_position, _ = read_uint16(rest)
    assert (vehicle_id, order_position) == (7, 2)


# --- OpenTTDClient scheduled dispatch edit commands ---

@pytest.mark.asyncio
async def test_client_set_scheduled_dispatch_payload():
    client = new_client()
    await client.set_scheduled_dispatch(7, True)
    parsed = decode_sent_command(client._protocol.sent[0])
    assert parsed["cmd"] == GameCommand.SchDispatch
    vehicle_id, rest = read_varuint(parsed["payload"])
    enabled, _ = read_uint8(rest)
    assert (vehicle_id, enabled) == (7, 1)

@pytest.mark.asyncio
async def test_client_add_dispatch_schedule_payload():
    client = new_client()
    await client.add_dispatch_schedule(7, -1234, 3000)
    parsed = decode_sent_command(client._protocol.sent[0])
    assert parsed["cmd"] == GameCommand.SchDispatchAddNewSchedule
    vehicle_id, rest = read_varuint(parsed["payload"])
    start_tick, rest = read_varuint_signed(rest)
    duration, _ = read_varuint(rest)
    assert (vehicle_id, start_tick, duration) == (7, -1234, 3000)

@pytest.mark.asyncio
async def test_client_remove_dispatch_schedule_payload():
    client = new_client()
    await client.remove_dispatch_schedule(7, 2)
    parsed = decode_sent_command(client._protocol.sent[0])
    assert parsed["cmd"] == GameCommand.SchDispatchRemoveSchedule
    vehicle_id, rest = read_varuint(parsed["payload"])
    schedule_index, _ = read_varuint(rest)
    assert (vehicle_id, schedule_index) == (7, 2)

@pytest.mark.asyncio
async def test_client_add_dispatch_slot_payload_defaults_and_extras():
    client = new_client()
    # Defaults: single slot, no interval/extra/flags/route.
    await client.add_dispatch_slot(7, 1, 500)
    # Bulk: three extra slots spaced 250 ticks apart, with flags and route id.
    await client.add_dispatch_slot(7, 1, 500, interval=250, extra_slots=3, slot_flags=5, route_id=2)

    def decode(payload):
        vehicle_id, rest = read_varuint(payload)
        schedule_index, rest = read_varuint(rest)
        offset, rest = read_varuint(rest)
        interval, rest = read_varuint(rest)
        extra_slots, rest = read_varuint(rest)
        slot_flags, rest = read_uint16(rest)
        route_id, _ = read_uint8(rest)
        return (vehicle_id, schedule_index, offset, interval, extra_slots, slot_flags, route_id)

    p0 = decode_sent_command(client._protocol.sent[0])
    p1 = decode_sent_command(client._protocol.sent[1])
    assert p0["cmd"] == GameCommand.SchDispatchAdd
    assert decode(p0["payload"]) == (7, 1, 500, 0, 0, 0, 0)
    assert decode(p1["payload"]) == (7, 1, 500, 250, 3, 5, 2)

@pytest.mark.asyncio
async def test_client_remove_dispatch_slot_payload():
    client = new_client()
    await client.remove_dispatch_slot(7, 1, 500)
    parsed = decode_sent_command(client._protocol.sent[0])
    assert parsed["cmd"] == GameCommand.SchDispatchRemove
    vehicle_id, rest = read_varuint(parsed["payload"])
    schedule_index, rest = read_varuint(rest)
    offset, _ = read_varuint(rest)
    assert (vehicle_id, schedule_index, offset) == (7, 1, 500)

@pytest.mark.asyncio
async def test_client_clear_dispatch_schedule_payload():
    client = new_client()
    await client.clear_dispatch_schedule(7, 1)
    parsed = decode_sent_command(client._protocol.sent[0])
    assert parsed["cmd"] == GameCommand.SchDispatchClear
    vehicle_id, rest = read_varuint(parsed["payload"])
    schedule_index, _ = read_varuint(rest)
    assert (vehicle_id, schedule_index) == (7, 1)

@pytest.mark.asyncio
async def test_client_set_dispatch_duration_payload():
    client = new_client()
    await client.set_dispatch_duration(7, 1, 4000)
    parsed = decode_sent_command(client._protocol.sent[0])
    assert parsed["cmd"] == GameCommand.SchDispatchSetDuration
    vehicle_id, rest = read_varuint(parsed["payload"])
    schedule_index, rest = read_varuint(rest)
    duration, _ = read_varuint(rest)
    assert (vehicle_id, schedule_index, duration) == (7, 1, 4000)

@pytest.mark.asyncio
async def test_client_set_dispatch_start_date_payload():
    client = new_client()
    await client.set_dispatch_start_date(7, 1, 1_000_000)
    parsed = decode_sent_command(client._protocol.sent[0])
    assert parsed["cmd"] == GameCommand.SchDispatchSetStartDate
    vehicle_id, rest = read_varuint(parsed["payload"])
    schedule_index, rest = read_varuint(rest)
    start_tick, _ = read_varuint_signed(rest)
    assert (vehicle_id, schedule_index, start_tick) == (7, 1, 1_000_000)


# --- OpenTTDClient.receive_ServerCommand dispatch ---

async def feed_command(client, cmd, payload):
    parsed = OpenTTDProtocol.receive_ServerCommand(None, memoryview(build_server_command_bytes(0, cmd, payload)))
    await client.receive_ServerCommand(None, **parsed)

@pytest.mark.asyncio
async def test_receive_change_timetable_updates_order_state():
    client = new_client()
    payload = bytearray()
    write_varuint(payload, 7)
    payload += (3).to_bytes(2, "little")
    payload.append(ModifyTimetableFlags.WaitTime)
    write_varuint(payload, 120)
    payload.append(0)
    await feed_command(client, GameCommand.ChangeTimetable, payload)
    assert client.get_vehicle_timetable(7) == {"orders": {3: {"wait_time": 120}}}

@pytest.mark.asyncio
async def test_receive_change_timetable_clear_field_sets_none():
    client = new_client()
    payload = bytearray()
    write_varuint(payload, 7)
    payload += (3).to_bytes(2, "little")
    payload.append(ModifyTimetableFlags.TravelTime)
    write_varuint(payload, 0)
    payload.append(ModifyTimetableCtrlFlag.ClearField)
    await feed_command(client, GameCommand.ChangeTimetable, payload)
    assert client.get_vehicle_timetable(7)["orders"][3]["travel_time"] is None

@pytest.mark.asyncio
async def test_receive_change_timetable_wait_fixed_stores_bool():
    client = new_client()
    payload = bytearray()
    write_varuint(payload, 7)
    payload += (0).to_bytes(2, "little")
    payload.append(ModifyTimetableFlags.SetWaitFixed)
    write_varuint(payload, 1)
    payload.append(0)
    await feed_command(client, GameCommand.ChangeTimetable, payload)
    assert client.get_vehicle_timetable(7)["orders"][0]["wait_time_fixed"] is True

@pytest.mark.asyncio
async def test_receive_autofill_timetable_updates_state():
    client = new_client()
    payload = bytearray()
    write_varuint(payload, 7)
    payload.append(1)
    payload.append(0)
    await feed_command(client, GameCommand.AutofillTimetable, payload)
    entry = client.get_vehicle_timetable(7)
    assert entry["autofill"] is True
    assert entry["autofill_preserve_wait_time"] is False

@pytest.mark.asyncio
async def test_receive_set_timetable_start_updates_state():
    client = new_client()
    payload = bytearray()
    write_varuint(payload, 7)
    payload.append(1)
    write_varuint_signed(payload, 555)
    await feed_command(client, GameCommand.SetTimetableStart, payload)
    entry = client.get_vehicle_timetable(7)
    assert entry["timetable_all"] is True
    assert entry["timetable_start"] == 555

@pytest.mark.asyncio
async def test_receive_set_vehicle_on_time_updates_state():
    client = new_client()
    payload = bytearray()
    write_varuint(payload, 7)
    payload.append(1)
    await feed_command(client, GameCommand.SetVehicleOnTime, payload)
    assert client.get_vehicle_timetable(7)["on_time_apply_to_group"] is True

@pytest.mark.asyncio
async def test_receive_unknown_command_is_ignored():
    client = new_client()
    await feed_command(client, 999, bytearray())
    assert client.vehicle_timetables == {}


# --- get_vehicle_timetable ---

@pytest.mark.asyncio
async def test_get_vehicle_timetable_known_and_unknown():
    client = new_client()
    assert client.get_vehicle_timetable(7) is None
    payload = bytearray()
    write_varuint(payload, 7)
    payload.append(1)
    await feed_command(client, GameCommand.SetVehicleOnTime, payload)
    assert client.get_vehicle_timetable(7) is not None
    assert client.get_vehicle_timetable(42) is None