wehub-resource-sync
535 lines
18 KiB
Python
535 lines
18 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from pathlib import Path
|
|
|
|
from agent.compiler import load_model_globals
|
|
from agent.tools.probe_model import ProbeModelTool
|
|
from agent.tools.probe_model.helpers import ProbeSession
|
|
|
|
|
|
def _write_probe_fixture(script_path: Path, *, include_object_model: bool = True) -> None:
|
|
lines = [
|
|
"from __future__ import annotations",
|
|
"",
|
|
"from pathlib import Path",
|
|
"",
|
|
"from sdk import (",
|
|
" ArticulatedObject,",
|
|
" ArticulationType,",
|
|
" Box,",
|
|
" Cylinder,",
|
|
" MotionLimits,",
|
|
" Origin,",
|
|
" TestContext,",
|
|
")",
|
|
"",
|
|
"HERE = Path(__file__).resolve().parent",
|
|
"model = ArticulatedObject(name='probe_fixture')",
|
|
"panel = model.part('panel')",
|
|
"panel.visual(Box((1.0, 0.6, 0.08)), origin=Origin(xyz=(0.0, 0.0, 0.04)), name='panel_body')",
|
|
"panel.visual(Cylinder(radius=0.06, length=0.04), origin=Origin(xyz=(0.0, 0.0, 0.10)), name='knob')",
|
|
"keys = model.part('keys')",
|
|
"keys.visual(Box((0.10, 0.20, 0.04)), origin=Origin(xyz=(-0.24, 0.0, 0.02)), name='key_1')",
|
|
"keys.visual(Box((0.10, 0.20, 0.04)), origin=Origin(xyz=(0.00, 0.0, 0.02)), name='key_2')",
|
|
"keys.visual(Box((0.10, 0.20, 0.04)), origin=Origin(xyz=(0.24, 0.0, 0.02)), name='key_3')",
|
|
"arm = model.part('arm')",
|
|
"arm.visual(Box((0.30, 0.08, 0.08)), origin=Origin(xyz=(0.15, 0.0, 0.04)), name='arm_body')",
|
|
"model.articulation(",
|
|
" 'panel_to_keys',",
|
|
" ArticulationType.FIXED,",
|
|
" parent=panel,",
|
|
" child=keys,",
|
|
" origin=Origin(xyz=(0.0, 0.0, 0.08)),",
|
|
")",
|
|
"model.articulation(",
|
|
" 'panel_to_arm',",
|
|
" ArticulationType.REVOLUTE,",
|
|
" parent=panel,",
|
|
" child=arm,",
|
|
" origin=Origin(xyz=(0.45, 0.0, 0.08)),",
|
|
" axis=(0.0, 0.0, 1.0),",
|
|
" motion_limits=MotionLimits(lower=-0.5, upper=0.5, effort=2.0, velocity=1.0),",
|
|
")",
|
|
"",
|
|
]
|
|
if include_object_model:
|
|
lines.extend(
|
|
[
|
|
"object_model = model",
|
|
"",
|
|
]
|
|
)
|
|
lines.extend(
|
|
[
|
|
"def run_tests():",
|
|
" ctx = TestContext(object_model, asset_root=HERE)",
|
|
" return ctx.report()",
|
|
]
|
|
)
|
|
script_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
|
|
|
|
|
|
def _write_nested_joint_probe_fixture(script_path: Path) -> None:
|
|
lines = [
|
|
"from __future__ import annotations",
|
|
"",
|
|
"from pathlib import Path",
|
|
"",
|
|
"from sdk import ArticulatedObject, ArticulationType, Box, Origin, TestContext",
|
|
"",
|
|
"HERE = Path(__file__).resolve().parent",
|
|
"model = ArticulatedObject(name='nested_probe_fixture')",
|
|
"base = model.part('base')",
|
|
"base.visual(Box((1.0, 1.0, 1.0)), origin=Origin(xyz=(0.0, 0.0, 0.5)), name='base_body')",
|
|
"arm = model.part('arm')",
|
|
"arm.visual(Box((1.0, 0.2, 0.2)), origin=Origin(xyz=(0.5, 0.0, 0.1)), name='arm_body')",
|
|
"forearm = model.part('forearm')",
|
|
"forearm.visual(Box((1.0, 0.2, 0.2)), origin=Origin(xyz=(0.5, 0.0, 0.1)), name='forearm_body')",
|
|
"model.articulation(",
|
|
" 'base_to_arm',",
|
|
" ArticulationType.REVOLUTE,",
|
|
" parent=base,",
|
|
" child=arm,",
|
|
" origin=Origin(xyz=(1.0, 0.0, 1.0)),",
|
|
" axis=(0.0, 0.0, 1.0),",
|
|
")",
|
|
"model.articulation(",
|
|
" 'arm_to_forearm',",
|
|
" ArticulationType.FIXED,",
|
|
" parent=arm,",
|
|
" child=forearm,",
|
|
" origin=Origin(xyz=(1.0, 0.0, 0.0)),",
|
|
")",
|
|
"object_model = model",
|
|
"",
|
|
"def run_tests():",
|
|
" ctx = TestContext(object_model, asset_root=HERE)",
|
|
" return ctx.report()",
|
|
]
|
|
script_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
|
|
|
|
|
|
def _write_rotated_cylinder_probe_fixture(script_path: Path) -> None:
|
|
lines = [
|
|
"from __future__ import annotations",
|
|
"",
|
|
"from pathlib import Path",
|
|
"import math",
|
|
"",
|
|
"from sdk import ArticulatedObject, Cylinder, Origin, TestContext",
|
|
"",
|
|
"HERE = Path(__file__).resolve().parent",
|
|
"model = ArticulatedObject(name='rotated_cylinder_probe_fixture')",
|
|
"body = model.part('body')",
|
|
"body.visual(",
|
|
" Cylinder(radius=0.1, length=0.6),",
|
|
" origin=Origin(rpy=(math.pi / 4.0, math.pi / 4.0, 0.0)),",
|
|
" name='angled_cylinder',",
|
|
")",
|
|
"object_model = model",
|
|
"",
|
|
"def run_tests():",
|
|
" ctx = TestContext(object_model, asset_root=HERE)",
|
|
" return ctx.report()",
|
|
]
|
|
script_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
|
|
|
|
|
|
async def _run_probe(
|
|
script_path: Path,
|
|
code: str,
|
|
*,
|
|
timeout_ms: int = 5000,
|
|
include_stdout: bool = False,
|
|
) -> dict[str, object]:
|
|
tool = ProbeModelTool(sdk_package="sdk")
|
|
invocation = await tool.build(
|
|
{
|
|
"file_path": str(script_path),
|
|
"code": code,
|
|
"timeout_ms": timeout_ms,
|
|
"include_stdout": include_stdout,
|
|
}
|
|
)
|
|
result = await invocation.execute()
|
|
assert result.error is None
|
|
assert isinstance(result.output, dict)
|
|
return result.output
|
|
|
|
|
|
def test_probe_model_returns_structured_measurements(tmp_path: Path) -> None:
|
|
script_path = tmp_path / "model.py"
|
|
_write_probe_fixture(script_path)
|
|
|
|
output = asyncio.run(
|
|
_run_probe(
|
|
script_path,
|
|
"\n".join(
|
|
[
|
|
"panel = part('panel')",
|
|
"knob = visual('panel', 'knob')",
|
|
"emit({",
|
|
" 'part_names': [name(item) for item in parts()],",
|
|
" 'knob_mount': mount_report(knob, panel),",
|
|
"})",
|
|
]
|
|
),
|
|
)
|
|
)
|
|
|
|
assert output["ok"] is True
|
|
assert "elapsed_ms" not in output
|
|
result = output["result"]
|
|
assert isinstance(result, dict)
|
|
assert result["part_names"] == ["panel", "keys", "arm"]
|
|
knob_mount = result["knob_mount"]
|
|
assert knob_mount["metric_kind"] == "mount_review"
|
|
assert knob_mount["child"] == {"part": "panel", "visual": "knob"}
|
|
assert knob_mount["parent"] == {"part": "panel"}
|
|
|
|
|
|
def test_probe_model_captures_stdout_when_requested(tmp_path: Path) -> None:
|
|
script_path = tmp_path / "model.py"
|
|
_write_probe_fixture(script_path)
|
|
|
|
output = asyncio.run(
|
|
_run_probe(
|
|
script_path,
|
|
"print('debug line'); emit({'ok': True})",
|
|
include_stdout=True,
|
|
)
|
|
)
|
|
|
|
assert output["ok"] is True
|
|
assert "elapsed_ms" not in output
|
|
assert "debug line" in output["stdout"]
|
|
|
|
|
|
def test_probe_model_missing_emit_fails_clearly(tmp_path: Path) -> None:
|
|
script_path = tmp_path / "model.py"
|
|
_write_probe_fixture(script_path)
|
|
|
|
output = asyncio.run(_run_probe(script_path, "part('panel')"))
|
|
|
|
assert output["ok"] is False
|
|
assert output["error"]["type"] == "emit_contract"
|
|
assert "emit(value) was not called" in output["error"]["message"]
|
|
|
|
|
|
def test_probe_model_double_emit_fails_clearly(tmp_path: Path) -> None:
|
|
script_path = tmp_path / "model.py"
|
|
_write_probe_fixture(script_path)
|
|
|
|
output = asyncio.run(_run_probe(script_path, "emit(1)\nemit(2)"))
|
|
|
|
assert output["ok"] is False
|
|
assert output["error"]["type"] == "emit_contract"
|
|
assert "exactly once" in output["error"]["message"]
|
|
|
|
|
|
def test_probe_model_snippet_exception_fails_clearly(tmp_path: Path) -> None:
|
|
script_path = tmp_path / "model.py"
|
|
_write_probe_fixture(script_path)
|
|
|
|
output = asyncio.run(_run_probe(script_path, "raise RuntimeError('boom')"))
|
|
|
|
assert output["ok"] is False
|
|
assert output["error"]["type"] == "snippet_exception"
|
|
assert "boom" in output["error"]["message"]
|
|
|
|
|
|
def test_probe_model_timeout_fails_clearly(tmp_path: Path) -> None:
|
|
script_path = tmp_path / "model.py"
|
|
_write_probe_fixture(script_path)
|
|
|
|
output = asyncio.run(
|
|
_run_probe(
|
|
script_path,
|
|
"import time\ntime.sleep(0.5)\nemit({'done': True})",
|
|
timeout_ms=100,
|
|
)
|
|
)
|
|
|
|
assert output["ok"] is False
|
|
assert output["error"]["type"] == "timeout"
|
|
|
|
|
|
def test_probe_model_runner_crash_returns_structured_error(tmp_path: Path) -> None:
|
|
script_path = tmp_path / "model.py"
|
|
_write_probe_fixture(script_path)
|
|
|
|
output = asyncio.run(_run_probe(script_path, "import os\nos._exit(7)"))
|
|
|
|
assert output["ok"] is False
|
|
assert output["error"]["type"] == "runner_process_error"
|
|
assert "code 7" in output["error"]["message"]
|
|
|
|
|
|
def test_probe_model_missing_object_model_fails_clearly(tmp_path: Path) -> None:
|
|
script_path = tmp_path / "model.py"
|
|
_write_probe_fixture(script_path, include_object_model=False)
|
|
|
|
output = asyncio.run(_run_probe(script_path, "emit({'ok': True})"))
|
|
|
|
assert output["ok"] is False
|
|
assert output["error"]["type"] == "load_failure"
|
|
assert "`object_model`" in output["error"]["message"]
|
|
|
|
|
|
def test_probe_model_lookup_failure_is_clear(tmp_path: Path) -> None:
|
|
script_path = tmp_path / "model.py"
|
|
_write_probe_fixture(script_path)
|
|
|
|
output = asyncio.run(_run_probe(script_path, "emit(summary(part('missing')))"))
|
|
|
|
assert output["ok"] is False
|
|
assert output["error"]["type"] == "lookup_failure"
|
|
assert "Unknown part" in output["error"]["message"]
|
|
|
|
|
|
def test_probe_model_non_serializable_result_is_clear(tmp_path: Path) -> None:
|
|
script_path = tmp_path / "model.py"
|
|
_write_probe_fixture(script_path)
|
|
|
|
output = asyncio.run(_run_probe(script_path, "emit({'bad': {1, 2, 3}})"))
|
|
|
|
assert output["ok"] is False
|
|
assert output["error"]["type"] == "non_serializable_result"
|
|
|
|
|
|
def test_probe_model_supports_layout_and_pose_sampling(tmp_path: Path) -> None:
|
|
script_path = tmp_path / "model.py"
|
|
_write_probe_fixture(script_path)
|
|
|
|
output = asyncio.run(
|
|
_run_probe(
|
|
script_path,
|
|
"\n".join(
|
|
[
|
|
"keys = [visual('keys', name) for name in ('key_1', 'key_2', 'key_3')]",
|
|
"emit({",
|
|
" 'layout': layout_report(keys, axis='x'),",
|
|
" 'poses': sample_poses(max_samples=4, seed=1),",
|
|
"})",
|
|
]
|
|
),
|
|
)
|
|
)
|
|
|
|
assert output["ok"] is True
|
|
result = output["result"]
|
|
layout = result["layout"]
|
|
assert layout["metric_kind"] == "layout_review"
|
|
assert layout["count"] == 3
|
|
assert len(layout["pitches"]) == 2
|
|
poses = result["poses"]
|
|
assert poses
|
|
assert isinstance(poses, list)
|
|
|
|
|
|
def test_probe_model_joint_position_is_world_space(tmp_path: Path) -> None:
|
|
script_path = tmp_path / "model.py"
|
|
_write_nested_joint_probe_fixture(script_path)
|
|
|
|
output = asyncio.run(
|
|
_run_probe(
|
|
script_path,
|
|
"\n".join(
|
|
[
|
|
"with pose(base_to_arm=1.57079632679):",
|
|
" emit({",
|
|
" 'joint_pos': position(joint('arm_to_forearm')),",
|
|
" 'child_pos': position(part('forearm')),",
|
|
" })",
|
|
]
|
|
),
|
|
)
|
|
)
|
|
|
|
assert output["ok"] is True
|
|
result = output["result"]
|
|
assert isinstance(result, dict)
|
|
joint_pos = result["joint_pos"]
|
|
child_pos = result["child_pos"]
|
|
assert isinstance(joint_pos, list)
|
|
assert isinstance(child_pos, list)
|
|
assert len(joint_pos) == 3
|
|
assert len(child_pos) == 3
|
|
for joint_value, child_value in zip(joint_pos, child_pos, strict=False):
|
|
assert abs(float(joint_value) - float(child_value)) < 1e-6
|
|
assert abs(float(joint_pos[0]) - 1.0) < 1e-6
|
|
assert abs(float(joint_pos[1]) - 1.0) < 1e-6
|
|
assert abs(float(joint_pos[2]) - 1.0) < 1e-6
|
|
|
|
|
|
def test_probe_model_non_aabb_measurements_use_exact_geometry(tmp_path: Path) -> None:
|
|
script_path = tmp_path / "model.py"
|
|
_write_rotated_cylinder_probe_fixture(script_path)
|
|
|
|
output = asyncio.run(
|
|
_run_probe(
|
|
script_path,
|
|
"\n".join(
|
|
[
|
|
"target = visual('body', 'angled_cylinder')",
|
|
"emit({",
|
|
" 'dims': dims(target),",
|
|
" 'center': center(target),",
|
|
" 'position': position(target),",
|
|
" 'projection': projection(target, 'xyz'),",
|
|
"})",
|
|
]
|
|
),
|
|
)
|
|
)
|
|
|
|
assert output["ok"] is True
|
|
result = output["result"]
|
|
dims = result["dims"]
|
|
center = result["center"]
|
|
position = result["position"]
|
|
intervals = result["projection"]["intervals"]
|
|
|
|
expected_dims = [float(intervals[axis][1]) - float(intervals[axis][0]) for axis in "xyz"]
|
|
expected_center = [
|
|
0.5 * (float(intervals[axis][0]) + float(intervals[axis][1])) for axis in "xyz"
|
|
]
|
|
|
|
assert isinstance(dims, list)
|
|
assert isinstance(center, list)
|
|
assert isinstance(position, list)
|
|
for observed, expected in zip(dims, expected_dims, strict=False):
|
|
assert abs(float(observed) - float(expected)) < 1e-6
|
|
for observed, expected in zip(center, expected_center, strict=False):
|
|
assert abs(float(observed) - float(expected)) < 1e-6
|
|
for observed, expected in zip(position, expected_center, strict=False):
|
|
assert abs(float(observed) - float(expected)) < 1e-6
|
|
|
|
|
|
def test_probe_session_reuses_projection_interval_cache(tmp_path: Path) -> None:
|
|
from sdk import TestContext
|
|
|
|
script_path = tmp_path / "model.py"
|
|
_write_probe_fixture(script_path)
|
|
globals_dict = load_model_globals(script_path, sdk_package="sdk")
|
|
object_model = globals_dict["object_model"]
|
|
ctx = TestContext(object_model, asset_root=tmp_path)
|
|
session = ProbeSession(object_model, ctx)
|
|
target = session.visual("panel", "knob")
|
|
|
|
calls: list[tuple[int, str]] = []
|
|
original = ctx._element_projection_interval
|
|
|
|
def counting_projection_interval(element: object, *, axis: str) -> tuple[float, float]:
|
|
calls.append((id(element), axis))
|
|
min_value, max_value = original(element, axis=axis)
|
|
return float(min_value), float(max_value)
|
|
|
|
ctx._element_projection_interval = counting_projection_interval # type: ignore[method-assign]
|
|
|
|
session.center(target)
|
|
session.dims(target)
|
|
session.projection(target, "xyz")
|
|
assert len(calls) == 3
|
|
element_id = calls[0][0]
|
|
assert calls == [(element_id, "x"), (element_id, "y"), (element_id, "z")]
|
|
|
|
session.center(target)
|
|
session.dims(target)
|
|
session.projection(target, "xyz")
|
|
assert len(calls) == 3
|
|
|
|
|
|
def _write_floating_probe_fixture(script_path: Path) -> None:
|
|
lines = [
|
|
"from __future__ import annotations",
|
|
"",
|
|
"from pathlib import Path",
|
|
"",
|
|
"from sdk import ArticulatedObject, ArticulationType, Box, Origin, TestContext",
|
|
"",
|
|
"HERE = Path(__file__).resolve().parent",
|
|
"model = ArticulatedObject(name='floating_probe_fixture')",
|
|
"base = model.part('base')",
|
|
"base.visual(Box((0.4, 0.4, 0.2)), origin=Origin(xyz=(0.0, 0.0, 0.1)), name='base_body')",
|
|
"payload = model.part('payload')",
|
|
"payload.visual(Box((0.2, 0.2, 0.2)), origin=Origin(xyz=(0.0, 0.0, 0.1)), name='payload_body')",
|
|
"model.articulation(",
|
|
" 'base_to_payload',",
|
|
" ArticulationType.FLOATING,",
|
|
" parent=base,",
|
|
" child=payload,",
|
|
" origin=Origin(xyz=(0.3, 0.0, 0.0)),",
|
|
" meta={'qc_samples': [Origin(xyz=(0.1, 0.0, 0.0))]},",
|
|
")",
|
|
"object_model = model",
|
|
"",
|
|
"def run_tests():",
|
|
" ctx = TestContext(object_model, asset_root=HERE)",
|
|
" return ctx.report()",
|
|
]
|
|
script_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
|
|
|
|
|
|
def test_probe_model_supports_floating_origin_pose_values(tmp_path: Path) -> None:
|
|
script_path = tmp_path / "model.py"
|
|
_write_floating_probe_fixture(script_path)
|
|
|
|
output = asyncio.run(
|
|
_run_probe(
|
|
script_path,
|
|
"\n".join(
|
|
[
|
|
"joint_obj = joint('base_to_payload')",
|
|
"before = position(part('payload'))",
|
|
"with pose(base_to_payload=Origin(xyz=(0.1, 0.2, 0.3))):",
|
|
" after = position(part('payload'))",
|
|
"emit({'summary': summary(joint_obj), 'before': before, 'after': after})",
|
|
]
|
|
),
|
|
)
|
|
)
|
|
|
|
assert output["ok"] is True
|
|
result = output["result"]
|
|
assert result["summary"]["type"] == "floating"
|
|
assert result["summary"]["motion"] == "rigid_6dof"
|
|
assert result["summary"]["pose_value_kind"] == "origin"
|
|
assert result["before"] == [0.3, 0.0, 0.0]
|
|
assert result["after"] == [0.4, 0.2, 0.3]
|
|
|
|
|
|
def test_probe_model_omits_null_placeholders_and_duplicate_mount_offsets(tmp_path: Path) -> None:
|
|
script_path = tmp_path / "model.py"
|
|
_write_probe_fixture(script_path)
|
|
|
|
output = asyncio.run(
|
|
_run_probe(
|
|
script_path,
|
|
"\n".join(
|
|
[
|
|
"panel = part('panel')",
|
|
"knob = visual('panel', 'knob')",
|
|
"emit({",
|
|
" 'part_summary': summary(panel),",
|
|
" 'visual_summary': summary(knob),",
|
|
" 'mount': mount_report(knob, panel),",
|
|
"})",
|
|
]
|
|
),
|
|
)
|
|
)
|
|
|
|
assert output["ok"] is True
|
|
result = output["result"]
|
|
part_summary = result["part_summary"]
|
|
visual_summary = result["visual_summary"]
|
|
mount = result["mount"]
|
|
|
|
assert "parent_part" not in part_summary
|
|
assert part_summary["visual_names"] == ["panel_body", "knob"]
|
|
assert visual_summary["parent_part"] == "panel"
|
|
assert "visual_names" not in visual_summary
|
|
assert "center_offset_xy" not in mount
|