156 lines
5.3 KiB
Python
156 lines
5.3 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
from experiments.scripts.output_test_bench import (
|
|
CaseSpec,
|
|
build_progress_snapshot,
|
|
build_benchmark_prompt,
|
|
discover_cases,
|
|
load_existing_results,
|
|
merge_artifacts,
|
|
scan_output_directory,
|
|
)
|
|
|
|
|
|
def test_discover_cases_extracts_case_metadata(tmp_path: Path) -> None:
|
|
dataset_root = tmp_path / "output_test"
|
|
case_dir = dataset_root / "simple" / "0000" / "00000797"
|
|
case_dir.mkdir(parents=True)
|
|
(case_dir / "description.txt").write_text("Create a part", encoding="utf-8")
|
|
(case_dir / "00000797.png").write_bytes(b"png")
|
|
(case_dir / "00000797_cip.py").write_text("# cip", encoding="utf-8")
|
|
(case_dir / "00000797_cq.py").write_text("# cq", encoding="utf-8")
|
|
(case_dir / "00000797.step").write_text("step", encoding="utf-8")
|
|
|
|
cases = discover_cases(dataset_root)
|
|
|
|
assert len(cases) == 1
|
|
case = cases[0]
|
|
assert case.case_id == "simple/0000/00000797"
|
|
assert case.case_slug == "simple__0000__00000797"
|
|
assert Path(case.description_path).name == "description.txt"
|
|
assert Path(case.image_path or "").name == "00000797.png"
|
|
assert Path(case.reference_cip_path or "").name == "00000797_cip.py"
|
|
assert Path(case.reference_cq_path or "").name == "00000797_cq.py"
|
|
assert Path(case.reference_step_path or "").name == "00000797.step"
|
|
|
|
|
|
def test_discover_cases_works_when_dataset_root_is_already_simple_dir(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
dataset_root = tmp_path / "simple"
|
|
case_dir = dataset_root / "0000" / "00000797"
|
|
case_dir.mkdir(parents=True)
|
|
(case_dir / "description.txt").write_text("Create a part", encoding="utf-8")
|
|
(case_dir / "00000797.png").write_bytes(b"png")
|
|
|
|
cases = discover_cases(dataset_root)
|
|
|
|
assert len(cases) == 1
|
|
case = cases[0]
|
|
assert case.case_id == "0000/00000797"
|
|
assert case.bucket_id == "0000"
|
|
assert case.sample_id == "00000797"
|
|
|
|
|
|
def test_build_benchmark_prompt_requires_full_automation(tmp_path: Path) -> None:
|
|
case = CaseSpec(
|
|
case_id="simple/0000/00000797",
|
|
case_slug="simple__0000__00000797",
|
|
bucket_id="0000",
|
|
sample_id="00000797",
|
|
case_dir=str(tmp_path),
|
|
description_path=str(tmp_path / "description.txt"),
|
|
image_path=str(tmp_path / "00000797.png"),
|
|
reference_cip_path=None,
|
|
reference_cq_path=None,
|
|
reference_step_path=None,
|
|
)
|
|
|
|
prompt = build_benchmark_prompt(
|
|
case=case,
|
|
description_text="Build a flange from the reference image.",
|
|
requested_output_dir=Path(
|
|
"/repo/workspace/experiments/output_test/r1/simple/0000/00000797"
|
|
),
|
|
execution_root=Path("/repo/workspace"),
|
|
)
|
|
|
|
assert "full user confirmation and authorization" in prompt
|
|
assert "Do not stop to ask for confirmation" in prompt
|
|
assert (
|
|
"target_output_dir: ./experiments/output_test/r1/simple/0000/00000797" in prompt
|
|
)
|
|
assert "<|code_file|>path/to/model.py</|code_file|>" in prompt
|
|
assert "Build a flange from the reference image." in prompt
|
|
|
|
|
|
def test_merge_artifacts_falls_back_to_directory_scan(tmp_path: Path) -> None:
|
|
output_dir = tmp_path / "workspace" / "case"
|
|
output_dir.mkdir(parents=True)
|
|
(output_dir / "model.py").write_text("print('ok')", encoding="utf-8")
|
|
(output_dir / "part.stl").write_text("solid", encoding="utf-8")
|
|
(output_dir / "part.step").write_text("step", encoding="utf-8")
|
|
|
|
scanned = scan_output_directory(output_dir)
|
|
merged = merge_artifacts(
|
|
tagged_artifacts={"code_path": None, "output_paths": []},
|
|
scanned_outputs=scanned,
|
|
requested_output_dir=output_dir,
|
|
)
|
|
|
|
assert merged["artifact_source"] == "dir_scan"
|
|
assert merged["code_path"] == output_dir / "model.py"
|
|
assert merged["stl_path"] == output_dir / "part.stl"
|
|
assert merged["step_path"] == output_dir / "part.step"
|
|
|
|
|
|
def test_load_existing_results_keeps_latest_per_case(tmp_path: Path) -> None:
|
|
case_results = tmp_path / "case_results.jsonl"
|
|
case_results.write_text(
|
|
"\n".join(
|
|
[
|
|
'{"case_id":"simple/0000/0001","status":"failed"}',
|
|
'{"case_id":"simple/0000/0002","status":"success"}',
|
|
'{"case_id":"simple/0000/0001","status":"success"}',
|
|
]
|
|
)
|
|
+ "\n",
|
|
encoding="utf-8",
|
|
)
|
|
|
|
loaded = load_existing_results(case_results)
|
|
|
|
assert sorted(loaded) == ["simple/0000/0001", "simple/0000/0002"]
|
|
assert loaded["simple/0000/0001"]["status"] == "success"
|
|
|
|
|
|
def test_build_progress_snapshot_counts_skipped_running_and_pending(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
results = [
|
|
{"case_id": "simple/0000/0001", "status": "success"},
|
|
{"case_id": "simple/0000/0002", "status": "failed"},
|
|
]
|
|
|
|
snapshot = build_progress_snapshot(
|
|
run_id="r1",
|
|
dataset_root=tmp_path,
|
|
results=results,
|
|
total_cases=5,
|
|
skipped_existing=2,
|
|
running_case_ids=["simple/0000/0003"],
|
|
pending_case_ids=["simple/0000/0004", "simple/0000/0005"],
|
|
status="running",
|
|
)
|
|
|
|
assert snapshot["run_id"] == "r1"
|
|
assert snapshot["total_cases"] == 5
|
|
assert snapshot["completed_cases"] == 2
|
|
assert snapshot["skipped_existing"] == 2
|
|
assert snapshot["running_cases"] == 1
|
|
assert snapshot["pending_cases"] == 2
|
|
assert snapshot["success"] == 1
|
|
assert snapshot["failed"] == 1
|