armcade

#!/usr/bin/env python
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
"""Stage-0 safety tests for the remote-control receiver (RemoteTargetServer) — no LiveKit, no arm.

Feeds real UDP datagrams to the server and asserts the translation into controller calls:
held-key jog -> command_delta, named-pose + freeze/estop edges, the latched e-stop, the deadman,
and rejection of non-finite / stale absolute targets. Run:  python test_remote_control.py
"""

from __future__ import annotations

import json
import socket
import sys
import time
from pathlib import Path

import numpy as np

from tomctl import control as cc

PORT = 19099
OK = True


def check(name: str, cond: bool) -> None:
    global OK
    print(f"  [{'PASS' if cond else 'FAIL'}] {name}")
    OK = OK and bool(cond)


class FakeController:
    """Captures the calls RemoteTargetServer.apply makes, with the attributes it touches."""

    def __init__(self) -> None:
        self.target_pos = np.array([0.1, 0.2, 0.3])
        self.commanded_pos = np.array([0.1, 0.2, 0.3])
        self.commanded_gripper = 0.0
        self.deltas: list[tuple] = []
        self.absolutes: list[tuple] = []
        self.calls: list[str] = []
        self.settling = False  # mimics a named-pose move in progress (go_home/go_rest)

    def command_delta(self, dx, dy, dz, dgrip):
        self.deltas.append((dx, dy, dz, dgrip))

    def command_absolute(self, x, y, z, g):
        self.absolutes.append((x, y, z, g))

    def is_settling(self):
        return self.settling

    def go_home(self):
        self.calls.append("home")
        self.settling = True

    def go_rest(self):
        self.calls.append("rest")
        self.settling = True


def send(sock, obj):
    sock.sendto(json.dumps(obj).encode(), ("127.0.0.1", PORT))
    time.sleep(0.06)  # let the receive thread parse before apply()


def main() -> int:
    torque: list[tuple] = []  # (arm, on) calls the receiver makes via the optional torque hook
    srv = cc.RemoteTargetServer(PORT, host="127.0.0.1", torque_cb=lambda arm, on: torque.append((arm, on)))
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    dt = 0.02
    try:
        # --- held-key jog -> command_delta, gated by deadman ---
        c = FakeController()
        send(sock, {"seq": 1, "type": "intent", "deadman": True, "held": ["up"]})
        srv.apply({"arm1": c}, dt)
        check("held 'up' with deadman -> command_delta", len(c.deltas) == 1 and c.deltas[0][0] > 0)

        c2 = FakeController()
        send(sock, {"seq": 2, "type": "intent", "deadman": False, "held": ["up"]})
        srv.apply({"arm1": c2}, dt)
        check("deadman False -> no jog", c2.deltas == [])

        c3 = FakeController()
        send(sock, {"seq": 3, "type": "intent", "deadman": True, "held": ["bogus"]})
        srv.apply({"arm1": c3}, dt)
        check("unknown held token ignored", c3.deltas == [])

        # --- session lifecycle edge, de-duped by id (only the bridge sends edges) ---
        c4 = FakeController()
        send(sock, {"seq": 4, "type": "intent", "deadman": True, "held": [], "edges": [{"id": 1, "action": "session_start"}]})
        srv.apply({"arm1": c4}, dt)
        srv.apply({"arm1": c4}, dt)  # same id must NOT fire twice
        send(sock, {"seq": 5, "type": "intent", "deadman": True, "held": [], "edges": [{"id": 1, "action": "session_start"}]})
        srv.apply({"arm1": c4}, dt)
        check("session_start fires exactly once (de-dup by id)", c4.calls.count("home") == 1)

        # --- e-stop is GONE: an "estop" edge is unknown now and simply ignored (jog still applies) ---
        ce = FakeController()
        send(sock, {"seq": 9, "type": "intent", "deadman": True, "held": ["up"], "edges": [{"id": 2, "action": "estop"}]})
        srv.apply({"arm1": ce}, dt)
        check("removed 'estop' edge is ignored -> jog still applies", len(ce.deltas) == 1 and ce.deltas[0][0] > 0)

        # --- absolute targets: reject non-finite, accept finite ---
        c8 = FakeController()
        send(sock, {"seq": 100, "targets": {"arm1": {"x": float("nan"), "y": 0.0, "z": 0.1, "gripper": 0.0}}})
        srv.apply({"arm1": c8}, dt)
        check("NaN absolute target rejected", c8.absolutes == [])
        send(sock, {"seq": 101, "targets": {"arm1": {"x": 0.2, "y": 0.1, "z": 0.15, "gripper": 50.0}}})
        srv.apply({"arm1": c8}, dt)
        check("finite absolute target applied", c8.absolutes and c8.absolutes[-1] == (0.2, 0.1, 0.15, 50.0))

        # --- deadman: silence freezes ---
        c9 = FakeController()
        c9.commanded_pos = np.array([5.0, 5.0, 5.0])
        srv.DEADMAN_S = 0.0  # force "silent"
        time.sleep(0.02)
        srv.apply({"arm1": c9}, dt)
        check("deadman timeout holds pose", np.allclose(c9.commanded_pos, c9.target_pos))

        # --- TWO ARMS: jog routed independently by the `arm` field ---
        srv.DEADMAN_S = 0.4  # restore (the deadman test above forced it to 0.0)
        a1, a2 = FakeController(), FakeController()
        send(sock, {"seq": 1, "type": "intent", "arm": "arm1", "deadman": True, "held": ["up"]})
        send(sock, {"seq": 1, "type": "intent", "arm": "arm2", "deadman": True, "held": ["down"]})
        srv.apply({"arm1": a1, "arm2": a2}, dt)
        check("arm1 jogs +dx, arm2 jogs -dx (independent routing)",
              len(a1.deltas) == 1 and a1.deltas[0][0] > 0 and len(a2.deltas) == 1 and a2.deltas[0][0] < 0)

        # session lifecycle edge dispatches to the NAMED arm only (arm2), never arm1
        d1, d2 = FakeController(), FakeController()
        send(sock, {"seq": 4, "type": "intent", "arm": "arm2", "deadman": True, "held": [],
                    "edges": [{"id": 3, "action": "session_start"}]})
        srv.apply({"arm1": d1, "arm2": d2}, dt)
        check("session_start routed to arm2 only (go_home on arm2, not arm1)",
              d2.calls.count("home") == 1 and "home" not in d1.calls)

        # missing `arm` field still defaults to arm1 (back-compat)
        e1 = FakeController()
        send(sock, {"seq": 50, "type": "intent", "deadman": True, "held": ["up"]})
        srv.apply({"arm1": e1}, dt)
        check("missing arm field defaults to arm1", len(e1.deltas) == 1 and e1.deltas[0][0] > 0)

        # --- remote session lifecycle: session_start = torque on + home ---
        torque.clear()
        s1 = FakeController()
        send(sock, {"seq": 60, "type": "intent", "arm": "arm1", "deadman": False, "held": [],
                    "edges": [{"id": 10, "action": "session_start"}]})
        srv.apply({"arm1": s1}, dt)
        check("session_start -> torque ON then go_home", torque == [("arm1", True)] and s1.calls == ["home"])

        # session_end = go_rest, then torque OFF only AFTER the rest move settles
        torque.clear()
        s1.calls.clear()
        send(sock, {"seq": 61, "type": "intent", "arm": "arm1", "deadman": False, "held": [],
                    "edges": [{"id": 11, "action": "session_end"}]})
        srv.apply({"arm1": s1}, dt)  # go_rest sets settling=True -> torque NOT released yet
        check("session_end -> go_rest, torque held while settling", s1.calls == ["rest"] and torque == [])
        s1.settling = False  # rest move completes
        send(sock, {"seq": 62, "type": "intent", "arm": "arm1", "deadman": False, "held": []})
        srv.apply({"arm1": s1}, dt)
        check("torque OFF once parked (deferred release)", torque == [("arm1", False)])

        # re-taking an arm before it finishes parking cancels the pending torque-off
        torque.clear()
        s2 = FakeController()
        send(sock, {"seq": 70, "type": "intent", "arm": "arm2", "deadman": False, "held": [],
                    "edges": [{"id": 20, "action": "session_end"}]})
        srv.apply({"arm2": s2}, dt)  # parking, release pending
        send(sock, {"seq": 71, "type": "intent", "arm": "arm2", "deadman": False, "held": [],
                    "edges": [{"id": 21, "action": "session_start"}]})
        srv.apply({"arm2": s2}, dt)  # re-taken -> torque ON, no later OFF
        s2.settling = False
        send(sock, {"seq": 72, "type": "intent", "arm": "arm2", "deadman": False, "held": []})
        srv.apply({"arm2": s2}, dt)
        check("re-take cancels the deferred torque-off", ("arm2", True) in torque and ("arm2", False) not in torque)
    finally:
        srv.close()
        sock.close()
    print("RECEIVER TESTS:", "PASS" if OK else "FAIL")
    return 0 if OK else 1


if __name__ == "__main__":
    raise SystemExit(main())