160 lines
4.9 KiB
Python
160 lines
4.9 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import json
|
|
from pathlib import Path
|
|
import struct
|
|
import sys
|
|
import time
|
|
from types import ModuleType
|
|
import unittest
|
|
|
|
|
|
class Dummy:
|
|
pass
|
|
|
|
|
|
def install_module(name: str, **attributes: object) -> ModuleType:
|
|
module = ModuleType(name)
|
|
for key, value in attributes.items():
|
|
setattr(module, key, value)
|
|
sys.modules[name] = module
|
|
return module
|
|
|
|
|
|
install_module("rclpy")
|
|
install_module("rclpy.node", Node=object)
|
|
for package, message_name in (
|
|
("brainco_hand_msgs.msg", "MotorStatus"),
|
|
("brainco_hand_msgs.msg", "SetMotorMulti"),
|
|
("diagnostic_msgs.msg", "DiagnosticStatus"),
|
|
("geometry_msgs.msg", "TwistStamped"),
|
|
("ros2_bridge_msgs.msg", "ArmStatus"),
|
|
("sensor_msgs.msg", "JointState"),
|
|
("std_srvs.srv", "Trigger"),
|
|
):
|
|
module = sys.modules.get(package) or install_module(package)
|
|
setattr(module, message_name, Dummy)
|
|
|
|
|
|
fake_omnisocket = install_module(
|
|
"omnisocket", CONTROL_DEFAULTS={}, MSG_TYPE_BINARY=2
|
|
)
|
|
|
|
module_path = Path(__file__).with_name("tg3_local_teleop.py")
|
|
spec = importlib.util.spec_from_file_location("tg3_local_teleop", module_path)
|
|
assert spec is not None and spec.loader is not None
|
|
teleop = importlib.util.module_from_spec(spec)
|
|
sys.modules[spec.name] = teleop
|
|
spec.loader.exec_module(teleop)
|
|
|
|
|
|
class FakeSession:
|
|
mode = "idle"
|
|
connect_count = 0
|
|
sequence = 0
|
|
sent_one = False
|
|
next_instance_id = 0
|
|
events: list[str] = []
|
|
|
|
def __init__(self) -> None:
|
|
type(self).next_instance_id += 1
|
|
self.instance_id = type(self).next_instance_id
|
|
|
|
def connect(self, **_kwargs: object) -> None:
|
|
type(self).connect_count += 1
|
|
type(self).events.append(f"connect:{self.instance_id}")
|
|
|
|
def stats(self) -> dict[str, int]:
|
|
return {"connected": 1, "registered": 1}
|
|
|
|
def recv(self, timeout_ms: int) -> tuple[str, int, bytes] | None:
|
|
if timeout_ms == 0:
|
|
return None
|
|
time.sleep(0.003)
|
|
if self.mode == "idle":
|
|
return None
|
|
if self.mode == "one_then_idle" and type(self).sent_one:
|
|
return None
|
|
type(self).sent_one = True
|
|
type(self).sequence += 1
|
|
data = {
|
|
"arm": {
|
|
"position": {"left": [0.0] * 7, "right": [0.0] * 7}
|
|
}
|
|
}
|
|
payload = json.dumps(data).encode()
|
|
packet = struct.pack(
|
|
"!4sQQI",
|
|
b"TG3A",
|
|
type(self).sequence,
|
|
time.time_ns(),
|
|
len(payload),
|
|
) + payload
|
|
return "expected-sender", 2, packet
|
|
|
|
def close(self) -> None:
|
|
type(self).events.append(f"close:{self.instance_id}")
|
|
return None
|
|
|
|
|
|
def config(refresh_s: float) -> dict[str, object]:
|
|
return {
|
|
"transport": "omnisocket",
|
|
"omnisocket_server": "127.0.0.1:14049",
|
|
"omnisocket_peer_id": "robot",
|
|
"omnisocket_expected_sender": "expected-sender",
|
|
"omnisocket_max_packet_age_ms": 300.0,
|
|
"omnisocket_idle_session_refresh_s": refresh_s,
|
|
}
|
|
|
|
|
|
class IdleSessionRefreshTest(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
FakeSession.connect_count = 0
|
|
FakeSession.sequence = 0
|
|
FakeSession.sent_one = False
|
|
FakeSession.next_instance_id = 0
|
|
FakeSession.events = []
|
|
fake_omnisocket.Session = FakeSession
|
|
|
|
def run_source(self, mode: str, run_s: float = 0.13) -> dict[str, object]:
|
|
FakeSession.mode = mode
|
|
source = teleop.LatestArmData(config(0.02))
|
|
source.start()
|
|
time.sleep(run_s)
|
|
source.close()
|
|
return source.metrics()
|
|
|
|
def test_idle_session_is_periodically_reconnected(self) -> None:
|
|
metrics = self.run_source("idle")
|
|
self.assertGreaterEqual(FakeSession.connect_count, 2)
|
|
self.assertGreaterEqual(metrics["idle_session_refreshes"], 1)
|
|
self.assertEqual(
|
|
FakeSession.events[:3],
|
|
["connect:1", "connect:2", "close:1"],
|
|
)
|
|
|
|
def test_continuous_valid_business_frames_prevent_refresh(self) -> None:
|
|
metrics = self.run_source("active", 0.08)
|
|
self.assertEqual(FakeSession.connect_count, 1)
|
|
self.assertEqual(metrics["idle_session_refreshes"], 0)
|
|
self.assertGreater(metrics["frames_accepted"], 1)
|
|
|
|
def test_session_refreshes_after_last_valid_frame(self) -> None:
|
|
metrics = self.run_source("one_then_idle")
|
|
self.assertGreaterEqual(FakeSession.connect_count, 2)
|
|
self.assertGreaterEqual(metrics["idle_session_refreshes"], 1)
|
|
self.assertEqual(metrics["frames_accepted"], 1)
|
|
|
|
def test_refresh_boundary_and_disable_switch(self) -> None:
|
|
due = teleop.LatestArmData._idle_session_refresh_due
|
|
self.assertFalse(due(10.0, 0.0, 0.0))
|
|
self.assertFalse(due(1.999, 0.0, 2.0))
|
|
self.assertTrue(due(2.0, 0.0, 2.0))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|