tests/test_robot_urdf.py

"""Synthetic, data-free tests for fpgm.robot.urdf.RobotModel.

No network access and no third-party assets: the URDF and its one mesh reference
are both written to ``tmp_path`` by the tests themselves, mirroring the shape of
the real PointWorld Franka Panda + Robotiq 2F-85 URDF (arm joints, a ``finger_joint``
gripper drive joint, mimic joints, and a collision-only ``*_sc``-style link) without
depending on it being present.
"""

from __future__ import annotations

import math
from pathlib import Path

import numpy as np
import pytest

from fpgm.robot.urdf import MimicJoint, RobotModel

_DRIVE_JOINT = "finger_joint"
_MIMIC_POS_JOINT = "test_mimic_pos_joint"
_MIMIC_NEG_JOINT = "test_mimic_neg_joint"

# A single-triangle ASCII STL, used only to exercise mesh loading + <mesh scale=...>.
# Unscaled, its vertices span x in [0, 1] and y in [0, 1].
_TRIANGLE_STL = """solid tri
facet normal 0 0 1
  outer loop
    vertex 0 0 0
    vertex 1 0 0
    vertex 0 1 0
  endloop
endfacet
endsolid tri
"""

_URDF_TEMPLATE = f"""<?xml version="1.0"?>
<robot name="synthetic_test">
  <link name="base_link"/>
  <link name="link1">
    <visual>
      <origin xyz="0 0 0.1" rpy="0 0 0"/>
      <geometry>
        <box size="0.05 0.05 0.05"/>
      </geometry>
    </visual>
  </link>
  <link name="link2">
    <visual>
      <origin xyz="0.5 0 0" rpy="0 0 0"/>
      <geometry>
        <mesh filename="tri.stl" scale="2 2 2"/>
      </geometry>
    </visual>
  </link>
  <link name="gripper_drive_link"/>
  <link name="mimic_pos_link"/>
  <link name="mimic_neg_link"/>
  <link name="link2_sc">
    <collision>
      <geometry>
        <box size="0.05 0.05 0.05"/>
      </geometry>
    </collision>
  </link>

  <joint name="joint1" type="revolute">
    <parent link="base_link"/>
    <child link="link1"/>
    <origin xyz="0 0 0" rpy="0 0 0"/>
    <axis xyz="0 0 1"/>
    <limit lower="-3.14159" upper="3.14159" effort="10" velocity="1"/>
  </joint>
  <joint name="joint2" type="revolute">
    <parent link="link1"/>
    <child link="link2"/>
    <origin xyz="1 0 0" rpy="0 0 0"/>
    <axis xyz="0 0 1"/>
    <limit lower="-3.14159" upper="3.14159" effort="10" velocity="1"/>
  </joint>
  <joint name="{_DRIVE_JOINT}" type="revolute">
    <parent link="link2"/>
    <child link="gripper_drive_link"/>
    <origin xyz="0 0 0" rpy="0 0 0"/>
    <axis xyz="1 0 0"/>
    <limit lower="0.0" upper="0.8" effort="10" velocity="1"/>
  </joint>
  <joint name="{_MIMIC_POS_JOINT}" type="revolute">
    <parent link="link2"/>
    <child link="mimic_pos_link"/>
    <origin xyz="0 0 0" rpy="0 0 0"/>
    <axis xyz="1 0 0"/>
    <limit lower="0.0" upper="0.8" effort="10" velocity="1"/>
    <mimic joint="{_DRIVE_JOINT}" multiplier="1" offset="0"/>
  </joint>
  <joint name="{_MIMIC_NEG_JOINT}" type="revolute">
    <parent link="link2"/>
    <child link="mimic_neg_link"/>
    <origin xyz="0 0 0" rpy="0 0 0"/>
    <axis xyz="1 0 0"/>
    <limit lower="-0.8" upper="0.0" effort="10" velocity="1"/>
    <mimic joint="{_DRIVE_JOINT}" multiplier="-1" offset="0"/>
  </joint>
  <joint name="link2_sc_joint" type="fixed">
    <parent link="link2"/>
    <child link="link2_sc"/>
    <origin xyz="0 0 0" rpy="0 0 0"/>
  </joint>
</robot>
"""


def _write_urdf(tmp_path: Path) -> Path:
    urdf_path = tmp_path / "synthetic.urdf"
    urdf_path.write_text(_URDF_TEMPLATE)
    (tmp_path / "tri.stl").write_text(_TRIANGLE_STL)
    return urdf_path


def _make_model(tmp_path: Path, **kwargs) -> RobotModel:
    urdf_path = _write_urdf(tmp_path)
    kwargs.setdefault("arm_joints", ["joint1", "joint2"])
    kwargs.setdefault("gripper_drive_joint", _DRIVE_JOINT)
    kwargs.setdefault("base_link", "base_link")
    return RobotModel(urdf_path, **kwargs)


class TestForwardKinematics:
    def test_link_pose_matches_hand_computed_rotation(self, tmp_path: Path):
        model = _make_model(tmp_path)

        # joint1 (base_link -> link1) rotates pi/2 about Z; joint2 (link1 -> link2)
        # is a fixed-at-zero offset of (1, 0, 0). A point at (1, 0, 0) rotated 90
        # degrees CCW about Z lands at (0, 1, 0) -- link2's origin in base_link.
        poses = model.link_poses(np.array([math.pi / 2, 0.0]), gripper=0.0)

        assert np.allclose(poses["link2"][:3, 3], [0.0, 1.0, 0.0], atol=1e-6)
        # link1's own origin is unaffected by its own joint's rotation.
        assert np.allclose(poses["link1"][:3, 3], [0.0, 0.0, 0.0], atol=1e-6)

        expected_rot = np.array([[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]])
        assert np.allclose(poses["link2"][:3, :3], expected_rot, atol=1e-6)

    def test_link_origins_matches_link_poses(self, tmp_path: Path):
        model = _make_model(tmp_path)
        q = np.array([0.2, -0.4])
        poses = model.link_poses(q, gripper=0.3)
        names, origins = model.link_origins(q, gripper=0.3)

        assert names == list(poses)
        for i, name in enumerate(names):
            assert np.allclose(origins[i], poses[name][:3, 3])


class TestMimicResolution:
    def test_mimic_joints_are_exposed_from_urdf_mimic_tags(self, tmp_path: Path):
        model = _make_model(tmp_path)
        by_name = {m.name: m for m in model.mimic_joints}

        assert set(by_name) == {_MIMIC_POS_JOINT, _MIMIC_NEG_JOINT}
        assert by_name[_MIMIC_POS_JOINT] == MimicJoint(
            name=_MIMIC_POS_JOINT, source=_DRIVE_JOINT, multiplier=1.0, offset=0.0
        )
        assert by_name[_MIMIC_NEG_JOINT] == MimicJoint(
            name=_MIMIC_NEG_JOINT, source=_DRIVE_JOINT, multiplier=-1.0, offset=0.0
        )

    def test_driving_joint_value_propagates_with_correct_sign(self, tmp_path: Path):
        model = _make_model(tmp_path)
        config = model.joint_config(np.array([0.0, 0.0]), gripper=0.5)

        lower, upper = model.gripper_limits
        q_drive = lower + 0.5 * (upper - lower)

        assert config[_DRIVE_JOINT] == pytest.approx(q_drive)
        assert config[_MIMIC_POS_JOINT] == pytest.approx(q_drive)
        assert config[_MIMIC_NEG_JOINT] == pytest.approx(-q_drive)


class TestGripperMapping:
    def test_gripper_zero_is_lower_limit(self, tmp_path: Path):
        model = _make_model(tmp_path)
        lower, _ = model.gripper_limits
        config = model.joint_config(np.array([0.0, 0.0]), gripper=0.0)
        assert config[_DRIVE_JOINT] == pytest.approx(lower)

    def test_gripper_one_is_upper_limit(self, tmp_path: Path):
        model = _make_model(tmp_path)
        _, upper = model.gripper_limits
        config = model.joint_config(np.array([0.0, 0.0]), gripper=1.0)
        assert config[_DRIVE_JOINT] == pytest.approx(upper)

    def test_gripper_half_is_midpoint(self, tmp_path: Path):
        model = _make_model(tmp_path)
        lower, upper = model.gripper_limits
        config = model.joint_config(np.array([0.0, 0.0]), gripper=0.5)
        assert config[_DRIVE_JOINT] == pytest.approx((lower + upper) / 2)

    @pytest.mark.parametrize("gripper,expected_frac", [(-5.0, 0.0), (2.5, 1.0)])
    def test_out_of_range_gripper_is_clipped(self, tmp_path: Path, gripper, expected_frac):
        model = _make_model(tmp_path)
        lower, upper = model.gripper_limits
        config = model.joint_config(np.array([0.0, 0.0]), gripper=gripper)
        assert config[_DRIVE_JOINT] == pytest.approx(lower + expected_frac * (upper - lower))


class TestValidation:
    def test_wrong_length_joint_positions_raises(self, tmp_path: Path):
        model = _make_model(tmp_path)
        with pytest.raises(ValueError, match="length"):
            model.joint_config(np.array([0.1, 0.2, 0.3]), gripper=0.0)

    def test_unknown_arm_joint_raises_with_available_list(self, tmp_path: Path):
        urdf_path = _write_urdf(tmp_path)
        with pytest.raises(ValueError, match="nonexistent_joint") as excinfo:
            RobotModel(
                urdf_path,
                base_link="base_link",
                arm_joints=["joint1", "nonexistent_joint"],
                gripper_drive_joint=_DRIVE_JOINT,
            )
        # the error should help debugging, not just say "not found"
        assert "joint1" in str(excinfo.value) or _DRIVE_JOINT in str(excinfo.value)

    def test_unknown_gripper_drive_joint_raises(self, tmp_path: Path):
        urdf_path = _write_urdf(tmp_path)
        with pytest.raises(ValueError, match="not_a_real_joint"):
            RobotModel(
                urdf_path,
                base_link="base_link",
                arm_joints=["joint1", "joint2"],
                gripper_drive_joint="not_a_real_joint",
            )

    def test_unknown_base_link_raises(self, tmp_path: Path):
        urdf_path = _write_urdf(tmp_path)
        with pytest.raises(ValueError, match="not_a_real_link"):
            RobotModel(
                urdf_path,
                base_link="not_a_real_link",
                arm_joints=["joint1", "joint2"],
                gripper_drive_joint=_DRIVE_JOINT,
            )


class TestPublicAccessors:
    """Accessors added so callers (e.g. fpgm.robot.kinematics/ik) never need to
    reach into RobotModel._urdf directly.
    """

    def test_arm_joints_and_base_link(self, tmp_path: Path):
        model = _make_model(tmp_path)
        assert model.arm_joints == ("joint1", "joint2")
        assert model.base_link == "base_link"

    def test_arm_joint_limits_matches_urdf(self, tmp_path: Path):
        model = _make_model(tmp_path)
        limits = model.arm_joint_limits()
        assert limits.shape == (2, 2)
        assert np.allclose(limits, [[-3.14159, 3.14159], [-3.14159, 3.14159]], atol=1e-4)

    def test_arm_joint_velocity_limits_matches_urdf(self, tmp_path: Path):
        model = _make_model(tmp_path)
        vels = model.arm_joint_velocity_limits()
        assert np.allclose(vels, [1.0, 1.0])

    def test_gripper_drive_velocity_limit(self, tmp_path: Path):
        model = _make_model(tmp_path)
        assert model.gripper_drive_velocity_limit == pytest.approx(1.0)

    def test_joint_origin_matches_urdf_xyz(self, tmp_path: Path):
        model = _make_model(tmp_path)
        origin = model.joint_origin("joint2")
        assert np.allclose(origin[:3, 3], [1.0, 0.0, 0.0], atol=1e-9)
        assert np.allclose(origin[:3, :3], np.eye(3), atol=1e-9)

    def test_joint_axis_matches_urdf(self, tmp_path: Path):
        model = _make_model(tmp_path)
        assert np.allclose(model.joint_axis("joint1"), [0.0, 0.0, 1.0])
        assert np.allclose(model.joint_axis(_DRIVE_JOINT), [1.0, 0.0, 0.0])

    def test_joint_type(self, tmp_path: Path):
        model = _make_model(tmp_path)
        assert model.joint_type("joint1") == "revolute"


class TestIntrospection:
    def test_link_and_joint_name_listings(self, tmp_path: Path):
        model = _make_model(tmp_path)
        assert set(model.link_names) == {
            "base_link",
            "link1",
            "link2",
            "gripper_drive_link",
            "mimic_pos_link",
            "mimic_neg_link",
            "link2_sc",
        }
        assert set(model.actuated_joint_names) == {"joint1", "joint2", _DRIVE_JOINT}


class TestVisualMeshes:
    def test_mesh_visual_gets_origin_and_scale_applied(self, tmp_path: Path):
        model = _make_model(tmp_path)
        meshes = model.visual_meshes()

        assert "link2" in meshes
        (mesh, transform), = meshes["link2"]

        # <mesh scale="2 2 2"/> on a unit-triangle-in-[0,1] should double its extent.
        assert np.allclose(mesh.vertices.max(axis=0)[:2], [2.0, 2.0], atol=1e-6)
        # <visual><origin xyz="0.5 0 0"/> should show up as the node's translation.
        assert np.allclose(transform[:3, 3], [0.5, 0.0, 0.0], atol=1e-6)
        assert np.allclose(transform[:3, :3], np.eye(3), atol=1e-6)

    def test_box_visual_is_returned_without_a_mesh_file(self, tmp_path: Path):
        model = _make_model(tmp_path)
        meshes = model.visual_meshes()

        assert "link1" in meshes
        (mesh, transform), = meshes["link1"]
        assert mesh.vertices.shape[0] > 0
        assert np.allclose(transform[:3, 3], [0.0, 0.0, 0.1], atol=1e-6)

    def test_links_without_visuals_are_absent(self, tmp_path: Path):
        model = _make_model(tmp_path)
        meshes = model.visual_meshes()
        assert "base_link" not in meshes
        assert "gripper_drive_link" not in meshes

    def test_collision_only_link_contributes_no_meshes(self, tmp_path: Path):
        # Mirrors the real PointWorld URDF's `*_sc` self-collision proxy links: a
        # link that has <collision> geometry but no <visual> at all must not show up
        # in visual_meshes(), since rendering it would draw collision capsules over
        # the real arm mesh.
        model = _make_model(tmp_path)
        meshes = model.visual_meshes()
        assert "link2_sc" not in meshes

    def test_load_meshes_false_skips_external_mesh_files_but_keeps_primitives(
        self, tmp_path: Path
    ):
        # <box>/<cylinder>/<sphere> are generated procedurally and don't need disk
        # I/O, so they still show up; only the external tri.stl visual is skipped.
        model = _make_model(tmp_path, load_meshes=False)
        meshes = model.visual_meshes()

        assert "link1" in meshes
        assert "link2" not in meshes