VLAarchtests/tests/test_compose_task_routed_proxy_summary.py

from __future__ import annotations

import json
from pathlib import Path

from eval.compose_task_routed_proxy_summary import _load_records


def test_load_records_can_select_specific_label(tmp_path: Path) -> None:
    path = tmp_path / "benchmark.json"
    payload = {
        "benchmark_config": {"controller": "model"},
        "iter8_bag": {
            "episode_records": [{"task_name": "bag", "success": 0.4}],
        },
        "iter9_bag": {
            "episode_records": [{"task_name": "bag", "success": 0.5}],
        },
    }
    path.write_text(json.dumps(payload), encoding="utf-8")

    label, records = _load_records(path, "iter9_bag")

    assert label == "iter9_bag"
    assert records == [{"task_name": "bag", "success": 0.5}]