From 31117fcf945552ec4c50e0b673f01ac371ccbb35 Mon Sep 17 00:00:00 2001
From: glannuzel <gaelle.lannuzel@gmail.com>
Date: Tue, 26 Aug 2025 16:04:45 +0200
Subject: [PATCH] Update test_reachy2.py

---
 tests/robots/test_reachy2.py | 84 +++++++++++++++++++++++++-----------
 1 file changed, 60 insertions(+), 24 deletions(-)

diff --git a/tests/robots/test_reachy2.py b/tests/robots/test_reachy2.py
index 5605d6bcb..a2c40c5bf 100644
--- a/tests/robots/test_reachy2.py
+++ b/tests/robots/test_reachy2.py
@@ -28,27 +28,27 @@ from lerobot.robots.reachy2 import (
 
 # {lerobot_keys: reachy2_sdk_keys}
 REACHY2_JOINTS = {
-    "head.neck.yaw": "neck_yaw.pos",
-    "head.neck.pitch": "neck_pitch.pos",
-    "head.neck.roll": "neck_roll.pos",
-    "head.l_antenna": "l_antenna.pos",
-    "head.r_antenna": "r_antenna.pos",
-    "r_arm.shoulder.pitch": "r_shoulder_pitch.pos",
-    "r_arm.shoulder.roll": "r_shoulder_roll.pos",
-    "r_arm.elbow.yaw": "r_elbow_yaw.pos",
-    "r_arm.elbow.pitch": "r_elbow_pitch.pos",
-    "r_arm.wrist.roll": "r_wrist_roll.pos",
-    "r_arm.wrist.pitch": "r_wrist_pitch.pos",
-    "r_arm.wrist.yaw": "r_wrist_yaw.pos",
-    "r_arm.gripper": "r_gripper.pos",
-    "l_arm.shoulder.pitch": "l_shoulder_pitch.pos",
-    "l_arm.shoulder.roll": "l_shoulder_roll.pos",
-    "l_arm.elbow.yaw": "l_elbow_yaw.pos",
-    "l_arm.elbow.pitch": "l_elbow_pitch.pos",
-    "l_arm.wrist.roll": "l_wrist_roll.pos",
-    "l_arm.wrist.pitch": "l_wrist_pitch.pos",
-    "l_arm.wrist.yaw": "l_wrist_yaw.pos",
-    "l_arm.gripper": "l_gripper.pos",
+    "neck_yaw.pos": "head.neck.yaw",
+    "neck_pitch.pos": "head.neck.pitch",
+    "neck_roll.pos": "head.neck.roll",
+    "l_antenna.pos": "head.l_antenna",
+    "r_antenna.pos": "head.r_antenna",
+    "r_shoulder_pitch.pos": "r_arm.shoulder.pitch",
+    "r_shoulder_roll.pos": "r_arm.shoulder.roll",
+    "r_elbow_yaw.pos": "r_arm.elbow.yaw",
+    "r_elbow_pitch.pos": "r_arm.elbow.pitch",
+    "r_wrist_roll.pos": "r_arm.wrist.roll",
+    "r_wrist_pitch.pos": "r_arm.wrist.pitch",
+    "r_wrist_yaw.pos": "r_arm.wrist.yaw",
+    "r_gripper.pos": "r_arm.gripper",
+    "l_shoulder_pitch.pos": "l_arm.shoulder.pitch",
+    "l_shoulder_roll.pos": "l_arm.shoulder.roll",
+    "l_elbow_yaw.pos": "l_arm.elbow.yaw",
+    "l_elbow_pitch.pos": "l_arm.elbow.pitch",
+    "l_wrist_roll.pos": "l_arm.wrist.roll",
+    "l_wrist_pitch.pos": "l_arm.wrist.pitch",
+    "l_wrist_yaw.pos": "l_arm.wrist.yaw",
+    "l_gripper.pos": "l_arm.gripper",
 }
 
 REACHY2_VEL = {
@@ -62,6 +62,8 @@ PARAMS = [
     {"with_mobile_base": False},
     {"with_mobile_base": False, "with_l_arm": False, "with_antennas": False},
     {"with_r_arm": False, "with_neck": False, "with_antennas": False},
+    {"use_external_commands": True},
+    {"use_external_commands": True, "with_mobile_base": False, "with_neck": False},
 ]
 
 
@@ -77,8 +79,8 @@ def _make_reachy2_sdk_mock():
 
     # Mock joints with some dummy positions
     joints = {
-        k: MagicMock(present_position=round(0.1 + 0.01 * i, 2))
-        for i, k in enumerate(REACHY2_JOINTS.keys())
+        k: MagicMock(present_position=i, goal_position=i + 0.1)
+        for i, k in enumerate(REACHY2_JOINTS.values())
     }
     r.joints = joints
 
@@ -158,10 +160,44 @@ def test_get_observation(reachy2):
     reachy2.connect()
     obs = reachy2.get_observation()
 
-    expected_keys = {f"{REACHY2_JOINTS[m]}" for m in reachy2.joints_dict.values()}
+    expected_keys = {m for m in reachy2.joints_dict.keys()}
     expected_keys.update(
         f"{v}" for v in REACHY2_VEL.keys() if reachy2.config.with_mobile_base
     )
     expected_keys.update(reachy2.cameras.keys())
     assert set(obs.keys()) == expected_keys
 
+    print(obs)
+
+    for motor in reachy2.joints_dict.keys():
+        assert (
+            obs[motor] == reachy2.reachy.joints[REACHY2_JOINTS[motor]].present_position
+        )
+    if reachy2.config.with_mobile_base:
+        for vel in REACHY2_VEL.keys():
+            assert obs[vel] == reachy2.reachy.mobile_base.odometry[REACHY2_VEL[vel]]
+
+
+def test_send_action(reachy2):
+    reachy2.connect()
+
+    action = {k: i * 10.0 for i, k in enumerate(reachy2.joints_dict.keys())}
+    if reachy2.config.with_mobile_base:
+        action.update({k: i * 0.1 for i, k in enumerate(REACHY2_VEL.keys(), start=1)})
+
+    returned = reachy2.send_action(action)
+
+    assert returned == action
+
+    if reachy2.config.with_mobile_base:
+        goal_speed = [i * 0.1 for i, _ in enumerate(REACHY2_VEL.keys(), start=1)]
+        reachy2.reachy.mobile_base.set_goal_speed.assert_called_once_with(*goal_speed)
+
+    if reachy2.config.use_external_commands:
+        reachy2.reachy.send_goal_positions.assert_not_called()
+        if reachy2.config.with_mobile_base:
+            reachy2.reachy.mobile_base.send_speed_command.assert_not_called()
+    else:
+        reachy2.reachy.send_goal_positions.assert_called_once()
+        if reachy2.config.with_mobile_base:
+            reachy2.reachy.mobile_base.send_speed_command.assert_called_once()