tests/test_maniskill_render.py

from __future__ import annotations

import json

import pytest

from dovla_cil.generation.maniskill_render import (
    _load_records,
    _rewrite_observation_refs,
    _validate_state_archive,
)


def test_offline_renderer_rewrites_grouped_observation_refs(tmp_path) -> None:
    shard = tmp_path / "shard_000000.jsonl"
    records = [
        {"record_id": "r0", "group_id": "g0", "observation_ref": None},
        {"record_id": "r1", "group_id": "g0", "observation_ref": None},
        {"record_id": "r2", "group_id": "g1", "observation_ref": None},
    ]
    shard.write_text(
        "".join(json.dumps(record) + "\n" for record in records),
        encoding="utf-8",
    )

    by_group, by_shard = _load_records([shard])
    _rewrite_observation_refs(
        by_shard,
        {
            "r0": ("observations.h5#initial_rgb_jpeg/0", "next/0"),
            "r1": ("observations.h5#initial_rgb_jpeg/0", "next/1"),
            "r2": ("observations.h5#initial_rgb_jpeg/1", "next/2"),
        },
    )

    rewritten = [json.loads(line) for line in shard.read_text().splitlines()]
    assert list(by_group) == ["g0", "g1"]
    assert [len(group) for group in by_group.values()] == [2, 1]
    assert rewritten[0]["observation_ref"] == rewritten[1]["observation_ref"]
    assert rewritten[2]["observation_ref"].endswith("/1")
    assert [record["next_observation_ref"] for record in rewritten] == [
        "next/0",
        "next/1",
        "next/2",
    ]


def test_offline_renderer_requires_versioned_before_after_states() -> None:
    _validate_state_archive({"version": 2, "initial": {}, "next": {}})

    with pytest.raises(ValueError, match="version 2"):
        _validate_state_archive({"group-id": {"actors": {}}})