Spaces:
Paused
Paused
| import json | |
| import os | |
| import stat | |
| import sys | |
| import textwrap | |
| from dataclasses import replace | |
| import pytest | |
| from ros_fixtures import INTERFACES, NODE_PY, as_file_blocks, cpp_package, python_package | |
| from api.codegen import ( | |
| DirectBackend, | |
| OpenCodeBackend, | |
| generate_code, | |
| mentioned_interfaces, | |
| nested_types, | |
| parse_file_blocks, | |
| retrieve, | |
| ) | |
| from api.llm import LLMError, ScriptedLLM | |
| from api.opencode_runner import OpenCodeWorkspace | |
| from api.ros_check import validate_package | |
| class TestValidatePackage: | |
| def test_valid_python_package(self): | |
| v = validate_package(python_package(), "humble", INTERFACES) | |
| assert v.valid, v.errors | |
| assert v.stats["packages"] == ["safety_stop"] | |
| assert v.stats["interfaces_used"] == ["geometry_msgs/msg/Twist", "sensor_msgs/msg/LaserScan"] | |
| assert v.warnings == [] | |
| def test_valid_cpp_package(self): | |
| v = validate_package(cpp_package(), "jazzy", INTERFACES) | |
| assert v.valid, v.errors | |
| assert v.stats["interfaces_used"] == ["std_msgs/msg/String"] | |
| def test_hallucinated_message_type(self): | |
| node = NODE_PY.replace("from sensor_msgs.msg import LaserScan", "from sensor_msgs.msg import Lidar") | |
| v = validate_package(python_package(node=node), "humble", INTERFACES) | |
| assert not v.valid | |
| assert any("sensor_msgs/msg/Lidar does not exist in ROS 2 humble" in e for e in v.errors) | |
| def test_hallucinated_cpp_include(self): | |
| files = cpp_package() | |
| files["talker_cpp/src/talker.cpp"] = files["talker_cpp/src/talker.cpp"].replace( | |
| "std_msgs/msg/string.hpp", "std_msgs/msg/text_line.hpp" | |
| ) | |
| v = validate_package(files, "humble", INTERFACES) | |
| assert any("std_msgs/msg/TextLine does not exist" in e for e in v.errors) | |
| def test_unindexed_package_is_a_warning(self): | |
| node = NODE_PY.replace("import rclpy\n", "import rclpy\nfrom vision_msgs.msg import Detection2DArray\n") | |
| v = validate_package(python_package(node=node), "humble", INTERFACES) | |
| assert v.valid and any( | |
| "vision_msgs/msg/Detection2DArray is from a package that is not indexed" in w for w in v.warnings | |
| ) | |
| def test_missing_dependency_warning(self): | |
| v = validate_package(python_package(deps=("rclpy",)), "humble", INTERFACES) | |
| assert v.valid | |
| assert any("uses geometry_msgs but does not declare it" in w for w in v.warnings) | |
| def test_syntax_error(self): | |
| v = validate_package(python_package(node="def broken(:\n"), "humble", INTERFACES) | |
| assert any("Python syntax error" in e for e in v.errors) | |
| def test_launch_without_generate_function(self): | |
| v = validate_package(python_package(launch="from launch import LaunchDescription\n"), "humble", INTERFACES) | |
| assert any("generate_launch_description" in e for e in v.errors) | |
| def test_bad_xml_launch(self): | |
| files = cpp_package() | |
| files["talker_cpp/launch/talker.launch.xml"] = "<launchx/>" | |
| assert any("root must be <launch>" in e for e in validate_package(files, "humble", INTERFACES).errors) | |
| def test_structure_errors(self): | |
| files = python_package() | |
| del files["safety_stop/setup.py"] | |
| assert any("without setup.py" in e for e in validate_package(files, "humble", INTERFACES).errors) | |
| files = cpp_package() | |
| files["talker_cpp/CMakeLists.txt"] = "project(x)\n" | |
| assert any("ament_package()" in e for e in validate_package(files, "humble", INTERFACES).errors) | |
| assert validate_package({"a.py": "x = 1"}, "humble", INTERFACES).errors == ["no package.xml found"] | |
| assert not validate_package({}, "humble", INTERFACES).valid | |
| def test_path_traversal(self, bad): | |
| v = validate_package({bad: "x", **python_package()}, "humble", INTERFACES) | |
| assert not v.valid and "must be relative" in v.errors[0] | |
| def test_custom_interfaces_in_package(self): | |
| files = python_package() | |
| files["safety_stop/msg/Obstacle.msg"] = "float32 distance\n" | |
| files["safety_stop/safety_stop/node.py"] += "\nfrom safety_stop.msg import Obstacle\n" | |
| assert validate_package(files, "humble", INTERFACES).valid | |
| def test_without_catalogue(self): | |
| v = validate_package(python_package(), "humble", {}) | |
| assert v.valid and any("no interface catalogue" in w for w in v.warnings) | |
| class TestRetrieval: | |
| def test_mentioned_interfaces(self): | |
| found = mentioned_interfaces("subscribe to LaserScan and publish geometry_msgs/Twist", INTERFACES["humble"]) | |
| assert found == ["geometry_msgs/msg/Twist", "sensor_msgs/msg/LaserScan"] | |
| assert mentioned_interfaces("nothing relevant here", INTERFACES["humble"]) == [] | |
| def test_nested_types(self): | |
| h = INTERFACES["humble"] | |
| assert nested_types(h["geometry_msgs/msg/Twist"], "geometry_msgs", h) == [ | |
| "geometry_msgs/msg/Vector3", | |
| "geometry_msgs/msg/Vector3", | |
| ] | |
| assert nested_types(h["sensor_msgs/msg/LaserScan"], "sensor_msgs", h) == ["std_msgs/msg/Header"] | |
| def test_retrieve_puts_exact_definitions_first(self, store): | |
| hits = retrieve( | |
| store, "stop when LaserScan shows an obstacle, publish Twist", "humble", "python", INTERFACES["humble"] | |
| ) | |
| sources = [h.source for h in hits] | |
| assert sources[:4] == [ | |
| "interface:humble:geometry_msgs/msg/Twist", | |
| "interface:humble:sensor_msgs/msg/LaserScan", | |
| "interface:humble:geometry_msgs/msg/Vector3", | |
| "interface:humble:std_msgs/msg/Header", | |
| ] | |
| assert "docs:pubsub" in sources | |
| def test_parse_file_blocks(): | |
| files = python_package() | |
| assert parse_file_blocks("Here you go:\n\n" + as_file_blocks(files)) == {k: v + "\n" for k, v in files.items()} | |
| def ros_settings(settings): | |
| settings.data_dir.mkdir(parents=True, exist_ok=True) | |
| settings.interfaces_file.write_text(json.dumps(INTERFACES)) | |
| from api.ros_check import load_interfaces | |
| load_interfaces.cache_clear() | |
| return settings | |
| class TestGenerateCodeDirect: | |
| def test_repair_loop(self, store, ros_settings): | |
| bad = python_package(node=NODE_PY.replace("import LaserScan", "import Lidar")) | |
| llm = ScriptedLLM([as_file_blocks(bad), as_file_blocks(python_package())]) | |
| result = generate_code( | |
| "stop the robot when LaserScan sees an obstacle, publish Twist", | |
| "humble", | |
| "python", | |
| DirectBackend(llm), | |
| store, | |
| ros_settings, | |
| ) | |
| assert result.valid and len(result.attempts) == 2 and result.backend == "direct" | |
| first_user = llm.calls[0][1].content | |
| assert "[ROS 2 humble Message sensor_msgs/msg/LaserScan]" in first_user # exact definition, grounded | |
| assert "sensor_msgs/msg/Lidar does not exist" in llm.calls[1][-1].content | |
| assert result.files["safety_stop/safety_stop/node.py"].strip() == NODE_PY.strip() | |
| def test_invalid_after_all_attempts(self, store, ros_settings): | |
| llm = ScriptedLLM(["no files", "still no files"]) | |
| result = generate_code( | |
| "a publisher of String messages", | |
| "humble", | |
| "python", | |
| DirectBackend(llm), | |
| store, | |
| ros_settings, | |
| max_attempts=2, | |
| ) | |
| assert not result.valid and result.validation.errors == ["no files were generated"] | |
| def test_bad_arguments(self, store, ros_settings): | |
| with pytest.raises(ValueError): | |
| generate_code("x" * 20, "rolling", "python", DirectBackend(ScriptedLLM([])), store, ros_settings) | |
| FAKE_OPENCODE = textwrap.dedent('''\ | |
| #!{python} | |
| """Stand-in for `opencode run`: writes a package into --dir, emits JSON events.""" | |
| import json, os, sys | |
| sys.path.insert(0, {tests!r}) | |
| from ros_fixtures import NODE_PY, python_package | |
| args = sys.argv[1:] | |
| ws = args[args.index("--dir") + 1] | |
| context = open(args[args.index("-f") + 1]).read() | |
| state = os.environ["XDG_DATA_HOME"] | |
| assert json.loads(os.environ["OPENCODE_CONFIG_CONTENT"])["permission"]["bash"] == "deny" | |
| log = os.path.join(state, "calls.jsonl") | |
| with open(log, "a") as f: | |
| f.write(json.dumps({{"argv": args, "context": context}}) + "\\n") | |
| if os.environ.get("FAKE_OPENCODE_FAIL"): | |
| print(json.dumps({{"type": "error", "error": {{"name": "APIError", "data": {{"message": "quota exceeded"}}}}}})) | |
| sys.exit(1) | |
| continuing = "--continue" in args | |
| node = NODE_PY if continuing else NODE_PY.replace("import LaserScan", "import Lidar") | |
| for path, content in python_package(node=node).items(): | |
| full = os.path.join(ws, path) | |
| os.makedirs(os.path.dirname(full), exist_ok=True) | |
| open(full, "w").write(content) | |
| print(json.dumps({{"type": "text", "part": {{"text": "fixed it" if continuing else "wrote the package"}}}})) | |
| ''') | |
| def fake_opencode(tmp_path, ros_settings): | |
| path = tmp_path / "opencode" | |
| path.write_text(FAKE_OPENCODE.format(python=sys.executable, tests=os.path.dirname(__file__))) | |
| path.chmod(path.stat().st_mode | stat.S_IEXEC) | |
| return replace(ros_settings, opencode_bin=str(path), codegen_backend="opencode") | |
| class TestOpenCodeBackend: | |
| def test_generate_then_repair_in_same_session(self, store, fake_opencode): | |
| with OpenCodeWorkspace(fake_opencode) as ws: | |
| result = generate_code( | |
| "stop the robot when LaserScan sees an obstacle", | |
| "humble", | |
| "python", | |
| OpenCodeBackend(ws), | |
| store, | |
| fake_opencode, | |
| ) | |
| calls = [json.loads(line) for line in (ws.state / "calls.jsonl").read_text().splitlines()] | |
| root = ws.root | |
| assert not root.exists() # workspace cleaned up | |
| assert result.valid and result.backend == "opencode" and len(result.attempts) == 2 | |
| assert "--continue" not in calls[0]["argv"] and "--continue" in calls[1]["argv"] | |
| assert "sensor_msgs/msg/LaserScan" in calls[0]["context"] | |
| assert "sensor_msgs/msg/Lidar does not exist" in calls[1]["argv"][1] | |
| assert result.attempts[1].transcript == "fixed it" | |
| assert "safety_stop/safety_stop/node.py" in result.files | |
| def test_opencode_error_surfaces(self, store, fake_opencode, monkeypatch): | |
| monkeypatch.setenv("FAKE_OPENCODE_FAIL", "1") | |
| with OpenCodeWorkspace(fake_opencode) as ws, pytest.raises(LLMError, match="quota exceeded"): | |
| generate_code( | |
| "a publisher of String messages", "humble", "python", OpenCodeBackend(ws), store, fake_opencode | |
| ) | |
| def test_missing_binary(self, store, ros_settings): | |
| s = replace(ros_settings, opencode_bin="/nonexistent/opencode") | |
| with OpenCodeWorkspace(s) as ws, pytest.raises(LLMError, match="not found"): | |
| generate_code("a publisher of String messages", "humble", "python", OpenCodeBackend(ws), store, s) | |