motian_mavbridge/tests/test_telemetry.py

import json
import math
import time
import unittest

from motian_bridge import Bridge, Config
from motian_bridge.dialects.v20 import VKFly
from motian_bridge.telemetry import Telemetry


def cfg(**kwargs):
    fields = dict(system_id="UAV000001", uav_type=7, serial_port="loop://",
                  serial_baudrate=57600, mqtt_host="localhost")
    fields.update(kwargs)
    return Config(**fields)


def sender(system=1, component=1):
    return VKFly.MAVLink(None, srcSystem=system, srcComponent=component)


def heartbeat(mav, armed=False, mode=4):
    return mav.heartbeat_encode(
        VKFly.MAV_TYPE_QUADROTOR, VKFly.MAV_AUTOPILOT_GENERIC,
        VKFly.MAV_MODE_FLAG_SAFETY_ARMED if armed else 0,
        mode, VKFly.MAV_STATE_ACTIVE).pack(mav)


def bms(mav, id_=0, voltage=22800, current=350, percent=76):
    return mav.vk_bms_status_encode(
        1000, voltage, current, 28, percent, id_, 3, 6,
        [0] * 30, 42, 95).pack(mav)


def position(mav):
    return mav.global_position_int_encode(
        1000, 312304000, 1214737000, 135500, 120500,
        520, 110, -30, 8730).pack(mav)


def decode_frame(frame):
    parser = VKFly.MAVLink(None)
    result = None
    for byte in frame:
        message = parser.parse_char(bytes((byte,)))
        if message is not None:
            result = message
    return result


class ParserTests(unittest.TestCase):
    def test_split_frames_crc_and_v1_heartbeat(self):
        telemetry = Telemetry(cfg())
        mav = sender()
        frame = heartbeat(mav, True)
        self.assertEqual(telemetry.feed(frame[:7], now=10), [])
        self.assertEqual(telemetry.feed(frame[7:] + position(mav), now=10),
                         [(274, True)])
        data = dict(telemetry.tick(now=10))[258]
        self.assertAlmostEqual(data["Latitude"], 31.2304)
        self.assertAlmostEqual(data["Altitude"], 120.5)
        self.assertEqual(data["VelocityZ"], -0.3)
        self.assertEqual(data["FlightMode"], "POSHOLD")

        corrupted = bytearray(position(mav))
        corrupted[-1] ^= 0xFF
        telemetry.feed(corrupted, now=11)
        # Corrupted position must not replace the good observation.
        self.assertEqual(telemetry.latest["GLOBAL_POSITION_INT"][1], 10)
        self.assertEqual(telemetry.feed(heartbeat(mav, False), now=12), [(274, False)])

        mav_v1 = sender()
        frame_v1 = mav_v1.heartbeat_encode(
            VKFly.MAV_TYPE_QUADROTOR, VKFly.MAV_AUTOPILOT_GENERIC, 128, 12,
            VKFly.MAV_STATE_ACTIVE).pack(mav_v1, force_mavlink1=True)
        self.assertEqual(frame_v1[0], 0xFE)
        self.assertEqual(Telemetry(cfg()).feed(frame_v1, now=1), [(274, True)])

    def test_flight_units_freshness_and_partial_state(self):
        telemetry = Telemetry(cfg())
        mav = sender()
        telemetry.feed(heartbeat(mav), now=0)
        initial = dict(telemetry.tick(now=0))[258]
        self.assertNotIn("Latitude", initial)
        self.assertEqual(initial["Armed"], False)
        telemetry.feed(mav.attitude_encode(
            1000, -0.5, 1.2, -math.pi / 2, 0, 0, 0).pack(mav), now=0.5)
        telemetry.feed(position(mav), now=0.5)
        telemetry.feed(mav.vfr_hud_encode(
            12.4, 13.1, 87, 50, 135.5, 0.8).pack(mav), now=0.5)
        telemetry.feed(mav.home_position_encode(
            312280000, 1214700000, 15000, 0, 0, 0, [1, 0, 0, 0],
            0, 0, 0).pack(mav), now=0.5)
        telemetry.feed(mav.gps_raw_int_encode(
            1000000, 3, 312304000, 1214737000, 135500,
            80, 110, 1300, 8700, 18).pack(mav), now=0.5)
        self.assertEqual(telemetry.tick(now=0.9), [])
        state = dict(telemetry.tick(now=1))[258]
        self.assertAlmostEqual(state["Roll"], math.degrees(-0.5), places=4)
        self.assertAlmostEqual(state["Pitch"], math.degrees(1.2), places=4)
        self.assertAlmostEqual(state["Yaw"], 270, places=4)
        self.assertAlmostEqual(state["GroundSpeed"], 13.1, places=5)
        self.assertEqual(state["HomeAltitude"], 15)
        self.assertGreater(state["SpaceDistToHome"], state["DistanceToHome"])
        self.assertEqual(state["Gps"]["SatelliteCount"], 18)
        self.assertEqual(state["Gps"]["LevelDivisor"], 0.8)
        self.assertNotIn("Temperature", state)
        self.assertNotIn("FlyTime", state)
        self.assertNotIn("IsMotorsOn", state)

        # Heartbeat is fresh, but all other observations have expired.
        telemetry.feed(heartbeat(mav), now=5.6)
        stale = dict(telemetry.tick(now=5.6))[258]
        self.assertNotIn("Latitude", stale)
        self.assertNotIn("Gps", stale)
        self.assertNotIn("HomeLatitude", stale)
        self.assertEqual(telemetry.tick(now=11), [])
        self.assertEqual(telemetry.feed(heartbeat(mav), now=11), [])
        self.assertEqual(dict(telemetry.tick(now=11))[258]["FlightMode"], "POSHOLD")

    def test_lock_to_first_flight_controller_and_reset(self):
        telemetry = Telemetry(cfg())
        gcs = sender(10, 190)
        foreign = sender(2, 1)
        selected = sender(1, 1)
        telemetry.feed(heartbeat(gcs), now=0)
        self.assertIsNone(telemetry.source)
        telemetry.feed(heartbeat(selected, True), now=0)
        telemetry.feed(heartbeat(foreign, False), now=0)
        telemetry.feed(position(foreign), now=0)
        self.assertEqual(telemetry.source, (1, 1))
        self.assertEqual(dict(telemetry.tick(now=0))[258]["Armed"], True)
        self.assertNotIn("Latitude", dict(telemetry.tick(now=1))[258])
        replacement = Telemetry(cfg())
        self.assertEqual(replacement.feed(heartbeat(foreign), now=2), [(274, False)])
        self.assertEqual(replacement.source, (2, 1))

    def test_bms_only_two_packs_rate_limit_and_stale(self):
        telemetry = Telemetry(cfg())
        mav = sender()
        telemetry.feed(heartbeat(mav), now=0)
        telemetry.feed(bms(mav, 0), now=0)
        single = dict(telemetry.tick(now=0))[257]
        self.assertEqual(single["Battery_VoltageLeft"], 22.8)
        self.assertEqual(single["Battery_Current"], 3.5)
        self.assertEqual(single["Battery_Temperature"], 28)
        self.assertEqual(single["cycleIndex"], 42)
        self.assertEqual(single["flags"], 3)
        telemetry.feed(bms(mav, 1, voltage=23000, current=-120, percent=50), now=0.2)
        self.assertEqual(telemetry.tick(now=0.2), [])
        dual = dict(telemetry.tick(now=1))[257]
        self.assertEqual(dual["Battery_VoltageRight"], 23)
        self.assertEqual(dual["Battery_CurrentRight"], -1.2)
        self.assertNotIn("Battery_Voltage", dual)
        self.assertNotIn("Battery_Remaining", dual)
        telemetry.feed(bms(mav, 7), now=2)
        self.assertNotIn(257, dict(telemetry.tick(now=2)))
        telemetry.feed(bms(mav, 1, percent=90), now=6)
        expired = dict(telemetry.tick(now=6))[257]
        self.assertNotIn("Battery_VoltageLeft", expired)
        self.assertEqual(expired["Battery_RemainingRight"], 90)

    def test_statustext_immediate_and_packs_without_heartbeat_ignored(self):
        telemetry = Telemetry(cfg())
        mav = sender()
        self.assertEqual(telemetry.feed(bms(mav), now=0), [])
        self.assertEqual(telemetry.tick(now=0), [])
        telemetry.feed(heartbeat(mav), now=1)
        status = mav.statustext_encode(4, "电池低".encode("utf-8")).pack(mav)
        self.assertEqual(telemetry.feed(status, now=1),
                         [(260, {"severity": 4, "msg": "电池低"})])

    def test_nonfinite_floats_are_omitted_from_json(self):
        telemetry = Telemetry(cfg())
        mav = sender()
        telemetry.feed(heartbeat(mav), now=0)
        telemetry.feed(mav.attitude_encode(
            1000, float("nan"), 0.25, float("inf"), 0, 0, 0).pack(mav), now=0)
        telemetry.feed(mav.vfr_hud_encode(
            float("nan"), 5, 0, 0, 0, float("inf")).pack(mav), now=0)
        report = dict(telemetry.tick(now=0))[258]
        self.assertNotIn("Roll", report)
        self.assertNotIn("Yaw", report)
        self.assertNotIn("AirSpeed", report)
        self.assertNotIn("ClimbRate", report)
        self.assertAlmostEqual(report["Pitch"], math.degrees(0.25), places=4)
        self.assertEqual(report["GroundSpeed"], 5)
        json.dumps(report, allow_nan=False)


class BridgeTelemetryTests(unittest.TestCase):
    def test_bridge_reports_and_raw_switch(self):
        from test_bridge import FakeMqtt, FakeSerial

        for raw in (True, False):
            with self.subTest(raw=raw):
                conf = cfg(publish_raw_mavlink=raw)
                mqtt = FakeMqtt()
                fake_serial = FakeSerial()
                mav = sender()
                fake_serial.to_read.append(heartbeat(mav, True) + position(mav) + bms(mav))
                bridge = Bridge(conf, mqtt, lambda *args, **kwargs: fake_serial)
                bridge.start()
                bridge._on_connect(mqtt, None, None, 0, None)
                try:
                    deadline = time.monotonic() + 2
                    while len(mqtt.published) < (4 if raw else 3) and time.monotonic() < deadline:
                        time.sleep(0.01)
                    by_type = {json.loads(payload)["DataType"]: (topic, json.loads(payload)["Data"])
                               for topic, payload, _, _ in mqtt.published}
                    self.assertEqual(526 in by_type, raw)
                    self.assertEqual(by_type[274], (conf.receive_topic, True))
                    self.assertEqual(by_type[257][0], conf.receive_topic)
                    self.assertEqual(by_type[258][1]["Latitude"], 31.2304)
                    self.assertEqual(by_type[258][0], conf.receive_topic)
                    # Turning off upstream raw frames does not affect #526 downlink.
                    bridge._on_message(mqtt, None, type("Command", (), {
                        "topic": conf.obtain_topic,
                        "retain": False,
                        "payload": json.dumps({
                            "DataType": 526, "UAVType": conf.uav_type,
                            "SystemID": conf.system_id, "TimeStamp": "2026-09-29T00:00:00+08:00",
                            "Data": "01FF",
                        }).encode(),
                    })())
                    self.assertEqual(fake_serial.written[-1], b"\x01\xff")
                finally:
                    bridge.stop()

    def test_serial_reconnect_rebinds_flight_controller(self):
        from test_bridge import FakeMqtt, FakeSerial

        class DisconnectedSerial(FakeSerial):
            def read(self, size):
                if self.to_read:
                    return self.to_read.pop(0)
                raise OSError("cable disconnected")

        first = DisconnectedSerial()
        first.to_read.append(heartbeat(sender(1), True))
        second = FakeSerial()
        second.to_read.append(heartbeat(sender(2), False))
        opened = []

        def factory(*args, **kwargs):
            opened.append(True)
            return first if len(opened) == 1 else second

        conf = cfg(publish_raw_mavlink=False, reconnect_seconds=0.01)
        mqtt = FakeMqtt()
        bridge = Bridge(conf, mqtt, factory)
        bridge._on_connect(mqtt, None, None, 0, None)
        bridge.start()
        try:
            deadline = time.monotonic() + 2
            while bridge._telemetry.source != (2, 1) and time.monotonic() < deadline:
                time.sleep(0.01)
            self.assertTrue(first.closed)
            self.assertEqual(bridge._telemetry.source, (2, 1))
            armed_events = [json.loads(payload)["Data"] for topic, payload, _, _ in mqtt.published
                            if json.loads(payload)["DataType"] == 274]
            self.assertEqual(armed_events, [True, False])
        finally:
            bridge.stop()


    def test_serial_loop_auto_downloads_parameters(self):
        from test_bridge import FakeMqtt, FakeSerial

        mav = sender()
        fake_serial = FakeSerial()
        fake_serial.to_read.append(heartbeat(mav, True))
        conf = cfg(publish_raw_mavlink=False)
        mqtt = FakeMqtt()
        bridge = Bridge(conf, mqtt, lambda *args, **kwargs: fake_serial)
        bridge._on_connect(mqtt, None, None, 0, None)
        bridge.start()
        try:
            deadline = time.monotonic() + 2
            while not fake_serial.written and time.monotonic() < deadline:
                time.sleep(0.01)
            types = [decode_frame(frame).get_type() for frame in fake_serial.written]
            self.assertIn("PARAM_REQUEST_LIST", types)
            request = next(decode_frame(frame) for frame in fake_serial.written
                           if decode_frame(frame).get_type() == "PARAM_REQUEST_LIST")
            self.assertEqual((request.target_system, request.target_component), (1, 1))
        finally:
            bridge.stop()


if __name__ == "__main__":
    unittest.main()