tests/test_robot_assets.py

"""Unit tests for scripts/fetch_robot_description.py's URDF-patching logic.

Pure XML string manipulation, deliberately: no network access happens on this
path, so these tests exercise :func:`patch_urdf_visuals` against a small
synthetic URDF rather than anything actually downloaded.
"""

from __future__ import annotations

import sys
import xml.etree.ElementTree as ET
from pathlib import Path

SCRIPTS_DIR = Path(__file__).resolve().parents[1] / "scripts"
if str(SCRIPTS_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPTS_DIR))

from fetch_robot_description import (  # noqa: E402
    COLLISION_MESH_SCALE,
    collect_referenced_meshes,
    patch_urdf_visuals,
)

# Mirrors the shape of the real panda_robotiq_85.urdf closely enough to exercise
# the patch: one Robotiq link whose visual is a missing .obj (the common case),
# one Robotiq link whose visual is the pad special case (already an .stl, but at
# the wrong -- pre-collision-rescale -- scale), and one arm link that must be
# left alone entirely.
_SYNTHETIC_URDF = """<?xml version="1.0"?>
<robot name="synthetic">
  <link name="robotiq_2f_85_base">
    <visual>
      <geometry>
        <mesh filename="meshes/robotiq-2f/visual/base.obj" scale="0.1 0.1 0.1"/>
      </geometry>
    </visual>
    <collision>
      <geometry>
        <mesh filename="meshes/robotiq-2f/collision/base.stl" scale="0.001 0.001 0.001"/>
      </geometry>
    </collision>
  </link>
  <link name="robotiq_2f_85_right_pad">
    <visual>
      <geometry>
        <mesh filename="meshes/robotiq-2f/visual/pad.stl" scale="0.0001 0.0001 0.0001"/>
      </geometry>
    </visual>
    <collision>
      <geometry>
        <mesh filename="meshes/robotiq-2f/collision/pad.stl" scale="0.001 0.001 0.001"/>
      </geometry>
    </collision>
  </link>
  <link name="panda_link0">
    <visual>
      <geometry>
        <mesh filename="meshes/visual/link0.dae"/>
      </geometry>
    </visual>
    <collision>
      <geometry>
        <mesh filename="meshes/collision/link0.stl"/>
      </geometry>
    </collision>
  </link>
</robot>
"""


def _mesh_filename_and_scale(link: ET.Element, tag: str) -> tuple[str, str | None]:
    mesh = link.find(f"{tag}/geometry/mesh")
    assert mesh is not None, f"no <{tag}>/geometry/mesh under {link.attrib.get('name')!r}"
    return mesh.attrib["filename"], mesh.attrib.get("scale")


class TestPatchUrdfVisuals:
    def test_obj_visual_repointed_to_collision_stl_with_corrected_scale(self) -> None:
        patched = ET.fromstring(patch_urdf_visuals(_SYNTHETIC_URDF))
        base = patched.find("./link[@name='robotiq_2f_85_base']")
        assert base is not None

        filename, scale = _mesh_filename_and_scale(base, "visual")
        assert filename == "meshes/robotiq-2f/collision/base.stl"
        assert scale == COLLISION_MESH_SCALE == "0.001 0.001 0.001"

    def test_collision_subtree_is_left_untouched(self) -> None:
        original = ET.fromstring(_SYNTHETIC_URDF)
        patched = ET.fromstring(patch_urdf_visuals(_SYNTHETIC_URDF))

        base_original = original.find("./link[@name='robotiq_2f_85_base']")
        base_patched = patched.find("./link[@name='robotiq_2f_85_base']")
        assert base_original is not None
        assert base_patched is not None

        before = _mesh_filename_and_scale(base_original, "collision")
        after = _mesh_filename_and_scale(base_patched, "collision")
        assert before == after == ("meshes/robotiq-2f/collision/base.stl", "0.001 0.001 0.001")

    def test_pad_special_case_maps_to_collision_stl_at_corrected_scale(self) -> None:
        patched = ET.fromstring(patch_urdf_visuals(_SYNTHETIC_URDF))
        pad = patched.find("./link[@name='robotiq_2f_85_right_pad']")
        assert pad is not None

        filename, scale = _mesh_filename_and_scale(pad, "visual")
        assert filename == "meshes/robotiq-2f/collision/pad.stl"
        # The original 0.0001 was calibrated for the missing visual mesh's own
        # units, not the collision STL's -- it must be corrected, not carried over.
        assert scale == "0.001 0.001 0.001"

    def test_non_robotiq_visual_is_left_untouched(self) -> None:
        patched = ET.fromstring(patch_urdf_visuals(_SYNTHETIC_URDF))
        arm_link = patched.find("./link[@name='panda_link0']")
        assert arm_link is not None

        filename, scale = _mesh_filename_and_scale(arm_link, "visual")
        assert filename == "meshes/visual/link0.dae"
        assert scale is None

    def test_patch_is_idempotent(self) -> None:
        once = patch_urdf_visuals(_SYNTHETIC_URDF)
        twice = patch_urdf_visuals(once)
        assert once == twice


# Mirrors the shape of the real pointworld franka_panda_robotiq_2f85.urdf closely
# enough to exercise mesh collection: a <visual> mesh, a <collision> mesh referencing
# a *different* file, a mesh filename repeated across two links (must de-duplicate),
# a nested relative path (must survive unnormalised), and a <box> primitive (must be
# ignored -- it has no filename attribute and isn't a <mesh> element at all).
_SYNTHETIC_MESH_URDF = """<?xml version="1.0"?>
<robot name="synthetic">
  <link name="panda_link0">
    <visual>
      <geometry>
        <mesh filename="meshes/visual/link0.dae"/>
      </geometry>
    </visual>
    <collision>
      <geometry>
        <mesh filename="meshes/collision/link0.stl"/>
      </geometry>
    </collision>
  </link>
  <link name="robotiq_85_base_link">
    <visual>
      <geometry>
        <mesh filename="meshes/visual/robotiq_85_base_link_fine.STL"/>
      </geometry>
    </visual>
    <collision>
      <geometry>
        <mesh filename="meshes/visual/robotiq_85_base_link_coarse.STL"/>
      </geometry>
    </collision>
  </link>
  <link name="left_outer_finger">
    <visual>
      <geometry>
        <mesh filename="meshes/visual/outer_finger_fine.STL"/>
      </geometry>
    </visual>
  </link>
  <link name="right_outer_finger">
    <visual>
      <!-- same file as left_outer_finger, referenced from a different link -->
      <mesh filename="meshes/visual/outer_finger_fine.STL"/>
    </visual>
  </link>
  <link name="some_box_link">
    <visual>
      <geometry>
        <box size="0.1 0.1 0.1"/>
      </geometry>
    </visual>
  </link>
</robot>
"""


class TestCollectReferencedMeshes:
    def test_returns_exactly_the_deduplicated_mesh_filenames(self) -> None:
        result = collect_referenced_meshes(_SYNTHETIC_MESH_URDF)
        assert result == {
            "meshes/visual/link0.dae",
            "meshes/collision/link0.stl",
            "meshes/visual/robotiq_85_base_link_fine.STL",
            "meshes/visual/robotiq_85_base_link_coarse.STL",
            "meshes/visual/outer_finger_fine.STL",
        }

    def test_box_primitive_is_ignored(self) -> None:
        result = collect_referenced_meshes(_SYNTHETIC_MESH_URDF)
        assert not any("box" in path for path in result)

    def test_nested_relative_paths_are_returned_unchanged(self) -> None:
        result = collect_referenced_meshes(_SYNTHETIC_MESH_URDF)
        # No normalisation: nested paths like "meshes/visual/link0.dae" must come
        # back byte-for-byte as written in the URDF, not resolved, collapsed, or
        # otherwise rewritten.
        assert "meshes/visual/link0.dae" in result
        assert "meshes/collision/link0.stl" in result