"""Keyboard composition uses the real LeRobot plugin; only the transport is fake.""" import math import runpy import subprocess import sys import unittest from pathlib import Path from types import SimpleNamespace from unittest.mock import Mock, patch from lerobot_robot_mujoco import LeKiwiSim, LeKiwiSimConfig from lerobot_robot_mujoco.lekiwi_sim import ARM, BASE, PROFILE from mujoco_control_bridge import RobotError from test_plugin import FakeClient SCRIPT = Path(__file__).resolve().parents[3] / "examples/lekiwi/teleoperate_sim.py" DEMO = runpy.run_path(str(SCRIPT)) KeyboardController = DEMO["KeyboardController"] ARM_KEYS = DEMO["ARM_KEYS"] class KeyboardTeleopTests(unittest.TestCase): def setUp(self): transport = patch("lerobot_robot_mujoco.lekiwi_sim.SimRobotClient", FakeClient) transport.start() self.addCleanup(transport.stop) self.robot = LeKiwiSim(LeKiwiSimConfig()) self.robot.connect() self.addCleanup(self.robot.disconnect) self.controller = KeyboardController(self.robot) def test_every_arm_key_moves_only_its_joint_in_degrees_or_percent(self): self.assertEqual(set(ARM_KEYS), set(ARM)) all_keys = [key for pair in ARM_KEYS.values() for key in pair] self.assertEqual(len(set(all_keys)), 12) self.assertFalse(set(all_keys) & set(self.robot.teleop_keys.values())) for channel, pair in ARM_KEYS.items(): for key, direction in zip(pair, (1, -1), strict=True): with self.subTest(channel=channel, key=key): self.robot.disconnect() self.robot.connect() controller = KeyboardController(self.robot, arm_speed=12, gripper_speed=30) initial = dict(controller.targets) result = controller.step([key], 0) for other in ARM: delta = direction * (1 if channel == "arm_gripper.pos" else 0.4) self.assertAlmostEqual( result[other], initial[other] + (delta if other == channel else 0) ) self.assertEqual({key: result[key] for key in BASE}, dict.fromkeys(BASE, 0)) def test_compose_all_six_arm_targets_and_base_in_one_action(self): initial = dict(self.controller.targets) result = self.controller.step("wauiotyv", 0) self.assertEqual(result["x.vel"], 0.1) self.assertEqual(result["y.vel"], 0.1) self.assertEqual(result["theta.vel"], 0) for channel in ARM: speed = 50 if channel == "arm_gripper.pos" else 20 self.assertAlmostEqual(result[channel], initial[channel] + speed / 30) self.assertEqual(len(self.robot._client.sent), 1) self.assertEqual(len(self.robot._client.sent[0]), 9) def test_opposing_keys_cancel_without_resetting_arm(self): initial = dict(self.controller.targets) keys = [key for pair in ARM_KEYS.values() for key in pair] + list("wsadzx") result = self.controller.step(keys, 0) for channel in ARM: self.assertEqual(result[channel], initial[channel]) self.assertEqual({key: result[key] for key in BASE}, dict.fromkeys(BASE, 0)) def test_expired_pulses_stop_base_and_hold_confirmed_not_measured_arm(self): first = self.controller.step("wu", 0) second = self.controller.step([], 0.1) self.assertGreater(second[ARM[0]], first[ARM[0]]) expired = self.controller.step([], 0.19) self.assertEqual(expired["x.vel"], 0) for channel in ARM: self.assertEqual(expired[channel], second[channel]) self.assertNotEqual(expired[ARM[0]], self.robot.get_observation()[ARM[0]]) self.assertEqual(self.controller.step([], 100), expired) # No catch-up jump. self.assertFalse(self.controller.active) def test_space_clears_motion_and_holds_pose_even_with_other_keys_in_batch(self): first = self.controller.step("wuv", 0) stopped = self.controller.step(" wur", 0.05) self.assertEqual(self.robot.speed_index, 0) for channel in ARM: self.assertEqual(stopped[channel], first[channel]) self.assertEqual(stopped["x.vel"], 0) self.assertEqual(self.controller.step([], 0.1), stopped) restarted = self.controller.step("u", 0.12) self.assertGreater(restarted[ARM[0]], stopped[ARM[0]]) def test_confirmed_clipping_prevents_joint_and_gripper_target_windup(self): for i in range(300): limited = self.controller.step("uv", i / 30) self.assertAlmostEqual(limited[ARM[0]], math.degrees(PROFILE["arm"][0]["max"])) self.assertEqual(limited["arm_gripper.pos"], 100) self.controller.step(" ", 10) reversed_action = self.controller.step("jb", 10.01) self.assertAlmostEqual(reversed_action[ARM[0]], limited[ARM[0]] - 20 / 30) self.assertAlmostEqual(reversed_action["arm_gripper.pos"], 100 - 50 / 30) for i in range(600): limited = self.controller.step("jb", 11 + i / 30) self.assertAlmostEqual(limited[ARM[0]], math.degrees(PROFILE["arm"][0]["min"])) self.assertEqual(limited["arm_gripper.pos"], 0) def test_base_speed_changes_once_per_input_event_not_each_pulse_tick(self): self.controller.step("r", 0) for now in (0.03, 0.06, 0.1, 0.17): self.controller.step([], now) self.assertEqual(self.robot.speed_index, 1) self.controller.step("f", 0.2) self.assertEqual(self.robot.speed_index, 0) def test_failed_send_does_not_commit_unconfirmed_targets(self): before = dict(self.controller.targets) with ( patch.object(self.robot, "send_action", side_effect=RobotError("TIMEOUT", "test")), self.assertRaises(RobotError), ): self.controller.step("u", 0) self.assertEqual(self.controller.targets, before) def test_main_always_releases_lease_and_restores_terminal(self): self.robot.disconnect() main = DEMO["main"] for outcome in (["q"], KeyboardInterrupt(), EOFError("test"), RobotError("STALE", "test")): with self.subTest(outcome=outcome): reader = Mock() if isinstance(outcome, BaseException): reader.side_effect = outcome else: reader.return_value = outcome with ( patch.dict(main.__globals__, {"read_keys": reader}), patch.dict( sys.modules, {"demo_control": SimpleNamespace(make_robot=lambda _: self.robot)}, ), patch.object(sys, "argv", [str(SCRIPT)]), patch.object(sys, "stdin") as stdin, patch("termios.tcgetattr", return_value=["saved"]), patch("termios.tcsetattr") as restore, patch("tty.setcbreak"), patch("builtins.print"), ): stdin.isatty.return_value = True stdin.fileno.return_value = 123 if isinstance(outcome, Exception): with self.assertRaises(type(outcome)): main() else: main() self.assertFalse(self.robot.is_connected) self.assertIsNone(self.robot._client) restore.assert_called_once_with(123, DEMO["termios"].TCSADRAIN, ["saved"]) def test_input_drain_is_bounded_and_eof_is_not_an_infinite_loop(self): with patch("select.select", return_value=([123], [], [])): with patch("os.read", return_value=b"u") as read: self.assertEqual(DEMO["read_keys"](123), ["u"] * 64) self.assertEqual(read.call_count, 64) with patch("os.read", return_value=b""), self.assertRaises(EOFError): DEMO["read_keys"](123) def test_invalid_speeds_rejected_before_loading_robot_or_connecting(self): for flag, value in ( ("--arm-speed", "nan"), ("--arm-speed", "91"), ("--gripper-speed", "inf"), ("--gripper-speed", "0"), ): with self.subTest(flag=flag, value=value): result = subprocess.run( [sys.executable, str(SCRIPT), flag, value], capture_output=True, text=True, timeout=5, ) self.assertEqual(result.returncode, 2) self.assertIn(f"{flag} 必须", result.stderr) if __name__ == "__main__": unittest.main()