File size: 1,731 Bytes
20c251e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | 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": {}}})
|