Consolidate tree traversal tests in the C and Python APIs.

PiperOrigin-RevId: 778883024
Change-Id: I0d30ed011e60a6da7148cebc1f715895fd0c6fd6
This commit is contained in:
Alessio Quaglino
2025-07-03 07:24:40 -07:00
committed by Copybara-Service
parent 465fb3234c
commit 69ce5dbd05
2 changed files with 62 additions and 74 deletions
+36 -48
View File
@@ -642,68 +642,56 @@ class SpecsTest(absltest.TestCase):
<mujoco> <mujoco>
<worldbody> <worldbody>
<body name="body1"> <body name="body1">
<site name="site1"/>
<body name="body3"> <body name="body3">
<site name="site1"/>
<site name="site2"/>
<site name="site3"/>
<site name="site4"/> <site name="site4"/>
<body name="body4">
<site name="site5"/>
<joint name="joint1"/>
<geom name="geom1" size="1"/>
</body>
</body> </body>
<geom name="geom1" size="1"/>
<geom name="geom2" size="1"/>
<site name="site2"/>
<site name="site3"/>
<geom name="geom3" size="1"/>
</body>
<body name="body2">
<site name="site5"/>
</body> </body>
<body name="body2"/>
</worldbody> </worldbody>
</mujoco> </mujoco>
""" """
spec = mujoco.MjSpec.from_string(main_xml) spec = mujoco.MjSpec.from_string(main_xml)
bodytype = mujoco.mjtObj.mjOBJ_BODY bodytype = mujoco.mjtObj.mjOBJ_BODY
self.assertLen(spec.bodies, 5) self.assertLen(spec.bodies, 4)
self.assertLen(spec.sites, 5) self.assertLen(spec.sites, 5)
self.assertLen(spec.worldbody.find_all('body'), 4) self.assertLen(spec.worldbody.find_all('body'), 3)
self.assertLen(spec.worldbody.find_all('site'), 5) self.assertLen(spec.worldbody.find_all('site'), 5)
self.assertLen(spec.worldbody.find_all('joint'), 1) self.assertLen(spec.worldbody.find_all('geom'), 3)
self.assertLen(spec.worldbody.find_all('geom'), 1)
self.assertEqual(spec.bodies[1].name, 'body1') self.assertEqual(spec.bodies[1].name, 'body1')
self.assertEqual(spec.bodies[2].name, 'body2') self.assertEqual(spec.bodies[2].name, 'body2')
self.assertEqual(spec.bodies[3].name, 'body3') self.assertEqual(spec.bodies[3].name, 'body3')
self.assertEqual(spec.bodies[4].name, 'body4') self.assertEqual(spec.bodies[1].parent, spec.bodies[0])
self.assertEqual(spec.bodies[1].parent, spec.worldbody) self.assertEqual(spec.bodies[2].parent, spec.bodies[0])
self.assertEqual(spec.bodies[2].parent, spec.worldbody)
self.assertEqual(spec.bodies[3].parent, spec.bodies[1]) self.assertEqual(spec.bodies[3].parent, spec.bodies[1])
self.assertEqual(spec.bodies[4].parent, spec.bodies[3]) self.assertLen(spec.worldbody.find_all(bodytype), 3)
self.assertLen(spec.worldbody.find_all(bodytype), 4) self.assertLen(spec.bodies[1].find_all(bodytype), 1)
self.assertLen(spec.bodies[1].find_all(bodytype), 2) self.assertEmpty(spec.bodies[3].find_all(bodytype))
self.assertLen(spec.bodies[3].find_all(bodytype), 1)
self.assertEqual(spec.worldbody.find_all('body')[0].name, 'body1')
self.assertEqual(spec.worldbody.find_all('body')[1].name, 'body2')
self.assertEqual(spec.worldbody.find_all('body')[2].name, 'body3')
self.assertEqual(spec.worldbody.find_all('body')[3].name, 'body4')
self.assertEqual(spec.bodies[1].find_all('body')[0].name, 'body3') self.assertEqual(spec.bodies[1].find_all('body')[0].name, 'body3')
self.assertEqual(spec.bodies[1].find_all('body')[1].name, 'body4') self.assertEmpty(spec.bodies[3].find_all('body'))
self.assertEqual(spec.bodies[3].find_all('body')[0].name, 'body4')
self.assertEmpty(spec.bodies[2].find_all('body')) self.assertEmpty(spec.bodies[2].find_all('body'))
self.assertEmpty(spec.bodies[4].find_all('body')) for i, body in enumerate(spec.worldbody.find_all('body')):
self.assertEqual(spec.worldbody.find_all('site')[0].name, 'site1') self.assertEqual(body.name, 'body' + str(i + 1))
self.assertEqual(spec.worldbody.find_all('site')[1].name, 'site2') for i, site in enumerate(spec.worldbody.find_all('site')):
self.assertEqual(spec.worldbody.find_all('site')[2].name, 'site3') self.assertEqual(site.name, 'site' + str(i + 1))
self.assertEqual(spec.worldbody.find_all('site')[3].name, 'site4') self.assertLen(spec.bodies[1].sites, 3)
self.assertEqual(spec.worldbody.find_all('site')[4].name, 'site5') self.assertLen(spec.bodies[2].sites, 1)
self.assertEmpty(spec.bodies[2].sites) self.assertLen(spec.bodies[3].sites, 1)
self.assertLen(spec.bodies[3].sites, 4) self.assertEqual(spec.bodies[1].sites[0].name, 'site1')
self.assertLen(spec.bodies[4].sites, 1) self.assertEqual(spec.bodies[1].sites[1].name, 'site2')
self.assertEqual(spec.bodies[3].sites[0].name, 'site1') self.assertEqual(spec.bodies[1].sites[2].name, 'site3')
self.assertEqual(spec.bodies[3].sites[1].name, 'site2') self.assertEqual(spec.bodies[3].sites[0].name, 'site4')
self.assertEqual(spec.bodies[3].sites[2].name, 'site3') self.assertEqual(spec.bodies[2].sites[0].name, 'site5')
self.assertEqual(spec.bodies[3].sites[3].name, 'site4') for body in spec.bodies:
self.assertEqual(spec.bodies[4].sites[0].name, 'site5') for site in body.sites:
self.assertEqual(spec.bodies[3].sites[0].parent, spec.bodies[3]) self.assertEqual(site.parent, body)
self.assertEqual(spec.bodies[3].sites[1].parent, spec.bodies[3])
self.assertEqual(spec.bodies[3].sites[2].parent, spec.bodies[3])
self.assertEqual(spec.bodies[3].sites[3].parent, spec.bodies[3])
self.assertEqual(spec.bodies[4].sites[0].parent, spec.bodies[4])
with self.assertRaises(ValueError) as cm: with self.assertRaises(ValueError) as cm:
spec.worldbody.find_all('actuator') spec.worldbody.find_all('actuator')
self.assertEqual( self.assertEqual(
@@ -711,9 +699,9 @@ class SpecsTest(absltest.TestCase):
'body.find_all supports the types: body, frame, geom, site,' 'body.find_all supports the types: body, frame, geom, site,'
' joint, light, camera.', ' joint, light, camera.',
) )
body4 = spec.worldbody.find_all('body')[3] body3 = spec.worldbody.find_all('body')[2]
body4.name = 'body4_new' body3.name = 'body3_new'
self.assertEqual(spec.bodies[4].name, 'body4_new') self.assertEqual(spec.bodies[3].name, 'body3_new')
def test_geom_list(self): def test_geom_list(self):
main_xml = """ main_xml = """
+26 -26
View File
@@ -72,18 +72,18 @@ TEST_F(MujocoTest, TreeTraversal) {
static constexpr char xml[] = R"( static constexpr char xml[] = R"(
<mujoco> <mujoco>
<worldbody> <worldbody>
<body name="body"> <body name="body1">
<body name="body1"> <site name="site1"/>
<body name="body2">
<site name="site4"/> <site name="site4"/>
</body> </body>
<site name="site1"/>
<geom name="geom1" size="1"/> <geom name="geom1" size="1"/>
<geom name="geom2" size="1"/> <geom name="geom2" size="1"/>
<site name="site2"/> <site name="site2"/>
<site name="site3"/> <site name="site3"/>
<geom name="geom3" size="1"/> <geom name="geom3" size="1"/>
</body> </body>
<body name="body2"> <body name="body3">
<site name="site5"/> <site name="site5"/>
</body> </body>
</worldbody> </worldbody>
@@ -95,9 +95,9 @@ TEST_F(MujocoTest, TreeTraversal) {
ASSERT_THAT(spec, NotNull()) << err.data(); ASSERT_THAT(spec, NotNull()) << err.data();
mjsBody* world = mjs_findBody(spec, "world"); mjsBody* world = mjs_findBody(spec, "world");
mjsBody* body = mjs_findBody(spec, "body");
mjsBody* body1 = mjs_findBody(spec, "body1"); mjsBody* body1 = mjs_findBody(spec, "body1");
mjsBody* body2 = mjs_findBody(spec, "body2"); mjsBody* body2 = mjs_findBody(spec, "body2");
mjsBody* body3 = mjs_findBody(spec, "body3");
mjsElement* site1 = mjs_findElement(spec, mjOBJ_SITE, "site1"); mjsElement* site1 = mjs_findElement(spec, mjOBJ_SITE, "site1");
mjsElement* site2 = mjs_findElement(spec, mjOBJ_SITE, "site2"); mjsElement* site2 = mjs_findElement(spec, mjOBJ_SITE, "site2");
mjsElement* site3 = mjs_findElement(spec, mjOBJ_SITE, "site3"); mjsElement* site3 = mjs_findElement(spec, mjOBJ_SITE, "site3");
@@ -110,30 +110,30 @@ TEST_F(MujocoTest, TreeTraversal) {
// test nonexistent // test nonexistent
EXPECT_EQ(mjs_firstElement(spec, mjOBJ_ACTUATOR), nullptr); EXPECT_EQ(mjs_firstElement(spec, mjOBJ_ACTUATOR), nullptr);
EXPECT_EQ(mjs_firstElement(spec, mjOBJ_LIGHT), nullptr); EXPECT_EQ(mjs_firstElement(spec, mjOBJ_LIGHT), nullptr);
EXPECT_EQ(mjs_firstChild(body, mjOBJ_CAMERA, /*recurse=*/true), nullptr); EXPECT_EQ(mjs_firstChild(body1, mjOBJ_CAMERA, /*recurse=*/true), nullptr);
EXPECT_EQ(mjs_firstChild(body, mjOBJ_TENDON, /*recurse=*/true), nullptr); EXPECT_EQ(mjs_firstChild(body1, mjOBJ_TENDON, /*recurse=*/true), nullptr);
// test first, nonrecursive // test first, nonrecursive
EXPECT_EQ(mjs_firstElement(spec, mjOBJ_BODY), world->element); EXPECT_EQ(mjs_firstElement(spec, mjOBJ_BODY), world->element);
EXPECT_EQ(site1, mjs_firstElement(spec, mjOBJ_SITE)); EXPECT_EQ(site1, mjs_firstElement(spec, mjOBJ_SITE));
EXPECT_EQ(site1, mjs_firstChild(body, mjOBJ_SITE, /*recurse=*/false)); EXPECT_EQ(site1, mjs_firstChild(body1, mjOBJ_SITE, /*recurse=*/false));
EXPECT_EQ(geom1, mjs_firstChild(body, mjOBJ_GEOM, /*recurse=*/false)); EXPECT_EQ(geom1, mjs_firstChild(body1, mjOBJ_GEOM, /*recurse=*/false));
EXPECT_EQ(site4, mjs_firstChild(body1, mjOBJ_SITE, /*recurse=*/false)); EXPECT_EQ(site4, mjs_firstChild(body2, mjOBJ_SITE, /*recurse=*/false));
EXPECT_EQ(site5, mjs_firstChild(body2, mjOBJ_SITE, /*recurse=*/false)); EXPECT_EQ(site5, mjs_firstChild(body3, mjOBJ_SITE, /*recurse=*/false));
// test first, recursive // test first, recursive
EXPECT_EQ(site1, mjs_firstChild(world, mjOBJ_SITE, /*recurse=*/true)); EXPECT_EQ(site1, mjs_firstChild(world, mjOBJ_SITE, /*recurse=*/true));
EXPECT_EQ(geom1, mjs_firstChild(world, mjOBJ_GEOM, /*recurse=*/true)); EXPECT_EQ(geom1, mjs_firstChild(world, mjOBJ_GEOM, /*recurse=*/true));
EXPECT_EQ(site4, mjs_firstChild(body1, mjOBJ_SITE, /*recurse=*/true)); EXPECT_EQ(site4, mjs_firstChild(body2, mjOBJ_SITE, /*recurse=*/true));
EXPECT_EQ(site5, mjs_firstChild(body2, mjOBJ_SITE, /*recurse=*/true)); EXPECT_EQ(site5, mjs_firstChild(body3, mjOBJ_SITE, /*recurse=*/true));
// text next, nonrecursive // text next, nonrecursive
EXPECT_EQ(site2, mjs_nextChild(body, site1, /*recursive=*/false)); EXPECT_EQ(site2, mjs_nextChild(body1, site1, /*recursive=*/false));
EXPECT_EQ(site3, mjs_nextChild(body, site2, /*recursive=*/false)); EXPECT_EQ(site3, mjs_nextChild(body1, site2, /*recursive=*/false));
EXPECT_EQ(nullptr, mjs_nextChild(body, site3, /*recursive=*/false)); EXPECT_EQ(nullptr, mjs_nextChild(body1, site3, /*recursive=*/false));
EXPECT_EQ(geom2, mjs_nextChild(body, geom1, /*recursive=*/false)); EXPECT_EQ(geom2, mjs_nextChild(body1, geom1, /*recursive=*/false));
EXPECT_EQ(geom3, mjs_nextChild(body, geom2, /*recursive=*/false)); EXPECT_EQ(geom3, mjs_nextChild(body1, geom2, /*recursive=*/false));
EXPECT_EQ(nullptr, mjs_nextChild(body, geom3, /*recursive=*/false)); EXPECT_EQ(nullptr, mjs_nextChild(body1, geom3, /*recursive=*/false));
EXPECT_EQ(mjs_nextElement(spec, site1), site2); EXPECT_EQ(mjs_nextElement(spec, site1), site2);
EXPECT_EQ(mjs_nextElement(spec, site2), site3); EXPECT_EQ(mjs_nextElement(spec, site2), site3);
EXPECT_EQ(mjs_nextElement(spec, site3), site4); EXPECT_EQ(mjs_nextElement(spec, site3), site4);
@@ -144,15 +144,15 @@ TEST_F(MujocoTest, TreeTraversal) {
EXPECT_EQ(mjs_nextElement(spec, geom3), nullptr); EXPECT_EQ(mjs_nextElement(spec, geom3), nullptr);
// text next, recursive // text next, recursive
EXPECT_EQ(site2, mjs_nextChild(body, site1, /*recursive=*/true)); EXPECT_EQ(site2, mjs_nextChild(body1, site1, /*recursive=*/true));
EXPECT_EQ(site3, mjs_nextChild(body, site2, /*recursive=*/true)); EXPECT_EQ(site3, mjs_nextChild(body1, site2, /*recursive=*/true));
EXPECT_EQ(site4, mjs_nextChild(body, site3, /*recursive=*/true)); EXPECT_EQ(site4, mjs_nextChild(body1, site3, /*recursive=*/true));
EXPECT_EQ(site4, mjs_nextChild(world, site3, /*recursive=*/true)); EXPECT_EQ(site4, mjs_nextChild(world, site3, /*recursive=*/true));
EXPECT_EQ(site5, mjs_nextChild(world, site4, /*recursive=*/true)); EXPECT_EQ(site5, mjs_nextChild(world, site4, /*recursive=*/true));
EXPECT_EQ(nullptr, mjs_nextChild(body, site5, /*recursive=*/true)); EXPECT_EQ(nullptr, mjs_nextChild(body1, site5, /*recursive=*/true));
EXPECT_EQ(geom2, mjs_nextChild(body, geom1, /*recursive=*/true)); EXPECT_EQ(geom2, mjs_nextChild(body1, geom1, /*recursive=*/true));
EXPECT_EQ(geom3, mjs_nextChild(body, geom2, /*recursive=*/true)); EXPECT_EQ(geom3, mjs_nextChild(body1, geom2, /*recursive=*/true));
EXPECT_EQ(nullptr, mjs_nextChild(body, geom3, /*recursive=*/true)); EXPECT_EQ(nullptr, mjs_nextChild(body1, geom3, /*recursive=*/true));
// check compilation ordering of sites // check compilation ordering of sites
mjModel* model = mj_compile(spec, nullptr); mjModel* model = mj_compile(spec, nullptr);