mcbebop/tests/test_sim.py

"""The simulator, checked with raw sockets rather than through our own client.

If the sim were only ever exercised by `arsdk/`, the two could agree about a
mistake and nothing would notice. So these tests do the handshake and the
framing by hand, and assert on bytes.
"""

import json
import socket
import struct
import time

import pytest

from mcbebop.arsdk.types import COMMAND_HEADER, BufferId, DataType, Frame
from mcbebop.protocol.types import Buffer
from mcbebop.sim import FakeBebop, _fallback_encode_args, _wire_value, load_specs, serve

SPECS = load_specs()


class Controller:
    """A hand-rolled controller: TCP handshake, then UDP both ways."""

    def __init__(self, sim: FakeBebop) -> None:
        self.udp = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.udp.bind(("127.0.0.1", 0))
        self.udp.settimeout(0.3)
        self.d2c_port = self.udp.getsockname()[1]
        request = {
            "d2c_port": self.d2c_port,
            "controller_type": "computer",
            "controller_name": "test",
            "arstream2_client_stream_port": 55004,
            "arstream2_client_control_port": 55005,
        }
        with socket.create_connection(("127.0.0.1", sim.discovery_port), timeout=2) as tcp:
            tcp.sendall(json.dumps(request).encode())
            raw = b""
            while b"\x00" not in raw:
                raw += tcp.recv(4096)
        self.reply = json.loads(raw.split(b"\x00")[0].decode())
        self.c2d = ("127.0.0.1", self.reply.get("c2d_port", 0))

    def send(self, data_type: int, buffer_id: int, payload: bytes, seq: int = 0) -> None:
        self.udp.sendto(Frame(DataType(data_type), buffer_id, seq, payload).encode(), self.c2d)

    def collect(self, seconds: float = 1.0) -> list[Frame]:
        frames: list[Frame] = []
        deadline = time.monotonic() + seconds
        while time.monotonic() < deadline:
            try:
                data, _ = self.udp.recvfrom(65535)
            except TimeoutError:
                continue
            frames.extend(Frame.decode_all(data))
        return frames

    def wait_for(self, predicate, seconds: float = 2.0) -> Frame:
        deadline = time.monotonic() + seconds
        while time.monotonic() < deadline:
            for frame in self.collect(0.3):
                if predicate(frame):
                    return frame
        pytest.fail("the sim never sent a matching frame")

    def close(self) -> None:
        self.udp.close()


@pytest.fixture
def sim():
    with FakeBebop() as fake:
        yield fake


@pytest.fixture
def controller(sim):
    client = Controller(sim)
    yield client
    client.close()


def events(frames: list[Frame]) -> list[tuple[tuple[int, int, int], bytes]]:
    return [
        (COMMAND_HEADER.unpack_from(f.payload), f.payload[COMMAND_HEADER.size :])
        for f in frames
        if f.buffer_id in (BufferId.D2C_ACK, BufferId.D2C_NON_ACK) and len(f.payload) >= COMMAND_HEADER.size
    ]


# -- the spec table ------------------------------------------------------
def test_every_command_is_indexed():
    assert len(SPECS) == 264


@pytest.mark.parametrize(
    ("name", "ids"),
    [
        # The table in bebop-2's protocol.md, which came off the aircraft.
        ("common.Common.AllStates", (0, 4, 0)),
        ("common.Settings.AllSettings", (0, 2, 0)),
        ("ardrone3.Piloting.TakeOff", (1, 0, 1)),
        ("ardrone3.Piloting.PCMD", (1, 0, 2)),
        ("ardrone3.MediaStreaming.VideoEnable", (1, 21, 0)),
    ],
)
def test_ids_match_the_observed_table(name, ids):
    assert SPECS[name].ids == ids


def test_class_names_are_reused_across_projects():
    # Why the index is keyed by full name: these are different commands.
    assert "common.SettingsState.ProductNameChanged" in SPECS
    assert "ardrone3.SettingsState.MotorErrorStateChanged" in SPECS
    assert SPECS["common.SettingsState.ProductNameChanged"].ids[0] == 0
    assert SPECS["ardrone3.SettingsState.MotorErrorStateChanged"].ids[0] == 1


def test_buffer_attribute_decides_the_frame_type():
    assert SPECS["ardrone3.Piloting.PCMD"].buffer == Buffer.NON_ACK
    assert SPECS["ardrone3.Piloting.Emergency"].buffer == Buffer.HIGH_PRIO
    assert SPECS["ardrone3.MediaStreaming.VideoEnable"].buffer == Buffer.ACK  # attribute absent


# -- encoding ------------------------------------------------------------
def test_strings_are_nul_terminated():
    spec = SPECS["common.SettingsState.ProductVersionChanged"]
    assert _fallback_encode_args(spec, {"software": "4.7.1", "hardware": "HW_05"}) == b"4.7.1\x00HW_05\x00"


def test_enums_travel_as_four_byte_positions():
    # pyparrot read these as one byte, which is wrong and shifts every
    # argument after them.
    spec = SPECS["ardrone3.PilotingState.FlyingStateChanged"]
    assert _fallback_encode_args(spec, {"state": "hovering"}) == struct.pack("<i", 2)
    assert _fallback_encode_args(spec, {"state": "landed"}) == struct.pack("<i", 0)


def test_mixed_widths_pack_little_endian():
    spec = SPECS["ardrone3.SettingsState.MotorFlightsStatusChanged"]
    raw = _fallback_encode_args(
        spec, {"nbFlights": 42, "lastFlightDuration": 611, "totalFlightDuration": 26_340}
    )
    assert raw == struct.pack("<HHI", 42, 611, 26_340)


def test_doubles_keep_the_no_position_sentinel_exactly():
    spec = SPECS["ardrone3.PilotingState.PositionChanged"]
    raw = _fallback_encode_args(spec, {"latitude": 500.0, "longitude": 500.0, "altitude": 500.0})
    assert struct.unpack("<ddd", raw) == (500.0, 500.0, 500.0)


def test_an_unknown_enum_member_is_an_error_not_a_zero():
    spec = SPECS["ardrone3.PilotingState.FlyingStateChanged"]
    with pytest.raises(KeyError):
        _wire_value(spec.args[0], "hovvering")


# -- handshake -----------------------------------------------------------
def test_handshake_is_accepted_and_names_a_c2d_port(sim, controller):
    assert controller.reply["status"] == 0
    assert controller.reply["c2d_port"] == sim.c2d_port
    assert controller.reply["arstream2_server_stream_port"] == 5004
    assert sim.handshakes[0]["controller_name"] == "test"


def test_a_second_controller_is_refused(sim, controller):
    second = Controller(sim)
    try:
        assert second.reply["status"] == 1
    finally:
        second.close()


def test_the_slot_is_free_again_after_release(sim, controller):
    assert sim.occupied
    sim.release()
    third = Controller(sim)
    try:
        assert third.reply["status"] == 0
    finally:
        third.close()


# -- what the drone sends ------------------------------------------------
def test_identity_burst_arrives_on_the_acknowledged_buffer(controller):
    frames = controller.collect(1.0)
    version = SPECS["common.SettingsState.ProductVersionChanged"].ids
    matching = [payload for ids, payload in events(frames) if ids == version]
    assert matching == [b"4.7.1\x00HW_05\x00"]
    assert all(
        f.data_type == DataType.DATA_WITH_ACK
        for f in frames
        if f.buffer_id == BufferId.D2C_ACK  # identity is sent acked, telemetry is not
    )


def test_one_sensor_self_test_deliberately_fails(controller):
    sensors = SPECS["common.CommonState.SensorsStatesListChanged"]
    names = sensors.args[0].members
    results = {}
    for ids, payload in events(controller.collect(1.0)):
        if ids == sensors.ids:
            index, state = struct.unpack("<iB", payload)
            results[names[index].name] = state
    assert len(results) == 6, results
    # A simulator where everything passes never exercises the code that
    # notices something has not.
    assert results["magnetometer"] == 0
    assert sum(results.values()) == 5


def test_telemetry_streams_at_roughly_five_hertz(controller):
    attitude = SPECS["ardrone3.PilotingState.AttitudeChanged"].ids
    count = sum(1 for ids, _ in events(controller.collect(1.2)) if ids == attitude)
    assert 3 <= count <= 9, f"{count} attitude events in 1.2s"


def test_the_drone_pings_us(controller):
    ping = controller.wait_for(lambda f: f.buffer_id == BufferId.PING)
    assert len(ping.payload) == 8  # the drone's own clock, to be echoed back
    assert ping.data_type == DataType.DATA


# -- what the drone receives ---------------------------------------------
def test_an_acknowledged_command_is_acknowledged(sim, controller):
    payload = COMMAND_HEADER.pack(1, 21, 0) + b"\x01"
    controller.send(DataType.DATA_WITH_ACK, BufferId.C2D_ACK, payload, seq=7)
    ack = controller.wait_for(lambda f: f.data_type == DataType.ACK)
    assert ack.buffer_id == BufferId.ack_for(BufferId.C2D_ACK) == 139
    assert ack.payload == b"\x07"  # the sequence acked, not the ack's own
    assert sim.received_ids().count((1, 21, 0)) == 1
    assert [r.args for r in sim.received if r.ids == (1, 21, 0)] == [b"\x01"]


def test_a_non_ack_command_is_recorded_but_not_acknowledged(sim, controller):
    controller.collect(0.4)  # drain the identity burst and its own frames
    controller.send(DataType.DATA, BufferId.C2D_NON_ACK, COMMAND_HEADER.pack(1, 0, 2) + bytes(9))
    deadline = time.monotonic() + 1.0
    while time.monotonic() < deadline and (1, 0, 2) not in sim.received_ids():
        time.sleep(0.02)
    assert (1, 0, 2) in sim.received_ids()
    assert not [f for f in controller.collect(0.4) if f.data_type == DataType.ACK]


def test_a_pong_is_counted_not_mistaken_for_a_command(sim, controller):
    ping = controller.wait_for(lambda f: f.buffer_id == BufferId.PING)
    before = len(sim.received)
    controller.send(DataType.DATA, BufferId.PONG, ping.payload)
    deadline = time.monotonic() + 1.0
    while time.monotonic() < deadline and sim.pongs == 0:
        time.sleep(0.02)
    assert sim.pongs >= 1
    # An 8-byte pong body parses as a command header and would otherwise be
    # filed as a command the controller never sent.
    assert len(sim.received) == before


def test_all_states_triggers_a_fresh_burst(sim, controller):
    controller.collect(0.6)
    version = SPECS["common.SettingsState.ProductVersionChanged"].ids
    controller.send(DataType.DATA_WITH_ACK, BufferId.C2D_ACK, COMMAND_HEADER.pack(0, 4, 0), seq=1)
    frames = controller.collect(1.0)
    assert any(ids == version for ids, _ in events(frames))
    assert (0, 4, 0) in sim.received_ids()


def test_a_truncated_datagram_does_not_stop_the_sim(sim, controller):
    # A size field larger than the datagram used to be an infinite loop.
    controller.udp.sendto(struct.pack("<BBBI", 4, 11, 1, 400) + b"short", controller.c2d)
    controller.send(DataType.DATA_WITH_ACK, BufferId.C2D_ACK, COMMAND_HEADER.pack(1, 21, 0) + b"\x00", 2)
    assert controller.wait_for(lambda f: f.data_type == DataType.ACK) is not None


# -- house rules ---------------------------------------------------------
def test_ports_are_ephemeral_by_default(sim):
    assert sim.discovery_port != 0 and sim.c2d_port != 0
    assert sim.discovery_port != sim.c2d_port


def test_serve_writes_nothing_to_stdout(capsys):
    # stdout is the MCP server's JSON-RPC transport and this module is
    # importable from it, so a stray print would corrupt the protocol.
    serve(0.3, discovery_port=0)
    assert capsys.readouterr().out == ""