Spaces:
Paused
Paused
| import pytest | |
| from conftest import GOOD_SCENE | |
| from api.llm import ScriptedLLM | |
| from api.mjcf_check import extract_mjcf, static_check, validate_mjcf | |
| from api.simulate import generate_scene | |
| def fenced(xml: str) -> str: | |
| return f"Here is the scene:\n```xml\n{xml}\n```\nNotes: crate is free." | |
| class TestExtract: | |
| def test_fenced_block(self): | |
| assert extract_mjcf(fenced(GOOD_SCENE)) == GOOD_SCENE | |
| def test_bare_document(self): | |
| assert extract_mjcf(f"Sure! {GOOD_SCENE} done") == GOOD_SCENE | |
| def test_skips_non_mjcf_blocks(self): | |
| text = "```python\nprint(1)\n```\n" + fenced(GOOD_SCENE) | |
| assert extract_mjcf(text) == GOOD_SCENE | |
| def test_nothing(self): | |
| assert extract_mjcf("I cannot do that.") is None | |
| class TestStaticCheck: | |
| def test_good_scene_passes(self): | |
| assert static_check(GOOD_SCENE, "unitree_g1") is None | |
| def test_rejections(self, xml, message): | |
| result = static_check(xml, "unitree_g1") | |
| assert result is not None and not result.valid and result.stage == "static" | |
| assert message in result.error | |
| def test_no_robot_forbids_includes(self): | |
| assert static_check("<mujoco><include file='g1.xml'/></mujoco>", None).error.startswith("<include") | |
| assert static_check("<mujoco><worldbody/></mujoco>", None) is None | |
| def test_size_limit(self): | |
| huge = "<mujoco><include file='g1.xml'/>" + "<!--" + "x" * 400_000 + "--></mujoco>" | |
| assert "larger than" in static_check(huge, "unitree_g1").error | |
| class TestLoadAndStep: | |
| def test_valid_scene(self, menagerie): | |
| r = validate_mjcf(GOOD_SCENE, "unitree_g1", menagerie, sim_seconds=0.2) | |
| assert r.valid and r.stage == "ok", r | |
| assert r.stats["nbody"] >= 4 and r.stats["simulated_seconds"] == pytest.approx(0.2, abs=0.01) | |
| def test_mujoco_compiler_error_is_reported(self, menagerie): | |
| r = validate_mjcf(GOOD_SCENE.replace('mass="3"', 'mass="3" frictoin="1"'), "unitree_g1", menagerie) | |
| assert not r.valid and r.stage == "load" | |
| assert "frictoin" in r.error | |
| def test_robot_root_body_required(self, menagerie): | |
| (menagerie / "unitree_g1" / "g1.xml").write_text("<mujoco><worldbody/></mujoco>") | |
| r = validate_mjcf(GOOD_SCENE, "unitree_g1", menagerie, sim_seconds=0.1) | |
| assert not r.valid and "pelvis" in r.error | |
| def test_missing_menagerie(self, tmp_path): | |
| with pytest.raises(FileNotFoundError, match="rag.ingest"): | |
| validate_mjcf(GOOD_SCENE, "unitree_g1", tmp_path / "nope") | |
| def test_timeout(self, menagerie): | |
| r = validate_mjcf(GOOD_SCENE, "unitree_g1", menagerie, sim_seconds=1e6, timeout=2) | |
| assert not r.valid and "longer than" in r.error | |
| class TestGenerateScene: | |
| def test_repair_loop_fixes_compiler_error(self, store, settings): | |
| bad = GOOD_SCENE.replace('mass="3"', 'mass="3" frictoin="1"') | |
| llm = ScriptedLLM([fenced(bad), fenced(GOOD_SCENE)]) | |
| result = generate_scene("G1 next to a slope and a crate", "unitree_g1", llm, store, settings) | |
| assert result.valid and len(result.attempts) == 2 | |
| assert result.mjcf == GOOD_SCENE | |
| repair_prompt = llm.calls[1][-1].content | |
| assert "frictoin" in repair_prompt and "stage 'load'" in repair_prompt | |
| assert llm.calls[1][-2].role == "assistant" | |
| def test_prompt_is_grounded(self, store, settings): | |
| llm = ScriptedLLM([fenced(GOOD_SCENE)]) | |
| result = generate_scene("a floor with friction and a box", "unitree_g1", llm, store, settings) | |
| system, user = llm.calls[0] | |
| assert '<include file="g1.xml"/>' in system.content | |
| assert "ROBOT FACTS" in user.content and "left_knee_joint" in user.content # read from the model file | |
| assert "REFERENCE CONTEXT" in user.content and "doc#geom" in user.content | |
| assert "doc#geom" in result.sources | |
| def test_gives_up_after_max_attempts(self, store, settings): | |
| llm = ScriptedLLM(["no xml here", "still none"]) | |
| result = generate_scene("anything at all", "unitree_g1", llm, store, settings, max_attempts=2) | |
| assert not result.valid and result.validation.stage == "extract" and len(result.attempts) == 2 | |
| def test_unknown_robot(self, store, settings): | |
| with pytest.raises(ValueError, match="unknown robot"): | |
| generate_scene("x" * 20, "atlas", ScriptedLLM([]), store, settings) | |
| def test_no_robot(self, store, settings): | |
| scene = "<mujoco><worldbody><geom name='floor' type='plane' size='0 0 1'/></worldbody></mujoco>" | |
| result = generate_scene("just a floor please", None, ScriptedLLM([fenced(scene)]), store, settings) | |
| assert result.valid | |
| SCHEMA_RST = """ | |
| .. dropdown:: :ref:`mujoco<mujoco>` :octicon:`star` | |
| .. grid-item:: | |
| :ref:`model<mujoco-model>` | |
| .. dropdown:: :ref:`(world)body<body>` :octicon:`sync` | |
| .. grid-item:: | |
| :ref:`name<body-name>` | |
| .. grid-item:: | |
| :ref:`pos<body-pos>` | |
| .. dropdown:: :ref:`geom<body-geom>` |*| | |
| .. grid-item:: | |
| :ref:`type<body-geom-type>` | |
| .. grid-item:: | |
| :ref:`friction<body-geom-friction>` | |
| """ | |
| def test_parse_mjcf_schema(): | |
| from rag.chunking import parse_mjcf_schema | |
| schema = parse_mjcf_schema(SCHEMA_RST) | |
| assert schema["body"] == {"element": "body", "path": "body", "attributes": ["name", "pos"]} | |
| assert schema["body-geom"]["attributes"] == ["type", "friction"] | |
| assert schema["mujoco"]["attributes"] == ["model"] | |
| def test_repair_prompt_lists_valid_attributes(store, settings): | |
| import json | |
| from rag.chunking import parse_mjcf_schema | |
| settings.data_dir.mkdir(parents=True, exist_ok=True) | |
| settings.mjcf_schema_file.write_text(json.dumps(parse_mjcf_schema(SCHEMA_RST))) | |
| (settings.menagerie_dir / "unitree_g1" / "scene.xml").write_text('<mujoco><include file="g1.xml"/></mujoco>') | |
| bad = GOOD_SCENE.replace('mass="3"', 'mass="3" frictoin="1"') | |
| llm = ScriptedLLM([fenced(bad), fenced(GOOD_SCENE)]) | |
| result = generate_scene("G1 next to a crate", "unitree_g1", llm, store, settings) | |
| first_prompt = llm.calls[0][1].content | |
| assert "MENAGERIE EXAMPLE" in first_prompt and "MJCF SCHEMA" in first_prompt | |
| assert "<geom> (body/geom): type, friction" in first_prompt | |
| assert "Valid attributes for that element:\n<geom> (body/geom): type, friction" in llm.calls[1][-1].content | |
| assert result.valid and result.sources[0] == "mujoco_menagerie:unitree_g1/scene.xml" | |
| def test_validator_ignores_rendering_settings(menagerie, monkeypatch): | |
| """Validation never renders, so a GL backend that is unusable here must not break it (seen in CI).""" | |
| monkeypatch.setenv("MUJOCO_GL", "not-a-backend") | |
| monkeypatch.setenv("PYOPENGL_PLATFORM", "not-a-backend") | |
| r = validate_mjcf(GOOD_SCENE, "unitree_g1", menagerie, sim_seconds=0.1) | |
| assert r.valid, r.error | |