tests/test_episode_spec.py

"""Tests for fpgm.datagen.episode_spec -- no GPU, no network, no real DROID data.

:class:`~fpgm.datagen.episode_spec.EpisodeMetadata` is a plain dataclass here (never
built via ``from_flows_h5``), so these tests exercise
:class:`~fpgm.datagen.episode_spec.EpisodeSpecResolver` purely against in-memory
metadata plus a :class:`~fpgm.config_datagen.DatagenProfile` pointed at a stub mesh
file this module creates itself.
"""

from __future__ import annotations

from pathlib import Path

import pytest

from fpgm.config_datagen import DatagenProfile, RoleConfig, RolesConfig, SceneAssetsConfig
from fpgm.datagen.episode_spec import (
    EpisodeMetadata,
    EpisodeSpec,
    EpisodeSpecResolver,
    ObjectRole,
    ObjectSpec,
    ShapeRepresentation,
    TaskNotParseableError,
)
from fpgm.types import ConfigError

_SCENE_ID = "8756300955"
_REAL_TASK = "put brick in drawer shelf and close drawer"
_UNPARSEABLE_TASK = "Do anything you like that takes multiple steps to complete."


def _stub_mesh(tmp_path: Path) -> Path:
    path = tmp_path / "manipulated.glb"
    path.write_bytes(b"stub-not-a-real-glb")  # resolver only checks existence, never loads it
    return path


def _profile(tmp_path: Path, *, with_parent_bridging: bool = False) -> DatagenProfile:
    profile = DatagenProfile()
    mesh_path = str(_stub_mesh(tmp_path))
    profile.scene_assets = SceneAssetsConfig(meshes={_SCENE_ID: {"manipulated": mesh_path}})
    if with_parent_bridging:
        profile.roles = RolesConfig(
            manipulated=RoleConfig(
                representation="mesh", plausible_scale_range=[0.03, 0.5], allow_scale=True,
                parent_role="fixture",
            ),
            fixture=RoleConfig(
                representation="observed_surface", plausible_scale_range=[0.5, 2.0],
                allow_scale=False,
            ),
        )
    return profile


def _metadata(task: str) -> EpisodeMetadata:
    return EpisodeMetadata(
        uuid="test-uuid", scene_id=_SCENE_ID, task=task,
        camera_serials={"ext1": "22008760", "ext2": "22008761"},
    )


class TestResolverDerivesRoles:
    def test_real_instruction_derives_brick_and_drawer_with_correct_roles(
        self, tmp_path: Path
    ) -> None:
        profile = _profile(tmp_path)
        spec = EpisodeSpecResolver(profile).resolve(_metadata(_REAL_TASK), camera_role="ext1")

        assert set(spec.labels) == {"brick", "drawer"}

        brick = spec.object("brick")
        assert brick.role is ObjectRole.MANIPULATED
        assert brick.representation is ShapeRepresentation.MESH
        assert brick.mesh_path is not None and brick.mesh_path.exists()
        assert brick.plausible_scale_range == (0.03, 0.5)
        assert brick.allow_scale is True

        drawer = spec.object("drawer")
        assert drawer.role is ObjectRole.FIXTURE
        assert drawer.representation is ShapeRepresentation.OBSERVED_SURFACE
        assert drawer.mesh_path is None
        assert drawer.plausible_scale_range == (0.5, 2.0)
        assert drawer.allow_scale is False

    def test_camera_role_and_serial_resolve_from_metadata(self, tmp_path: Path) -> None:
        profile = _profile(tmp_path)
        spec = EpisodeSpecResolver(profile).resolve(_metadata(_REAL_TASK), camera_role="ext2")
        assert spec.camera_role == "ext2"
        assert spec.camera_serial == "22008761"

    def test_parent_bridging_names_the_fixture_as_the_manipulated_objects_parent(
        self, tmp_path: Path
    ) -> None:
        profile = _profile(tmp_path, with_parent_bridging=True)
        spec = EpisodeSpecResolver(profile).resolve(_metadata(_REAL_TASK), camera_role="ext1")
        assert spec.parents == {"brick": "drawer"}
        assert spec.object("brick").parent_label == "drawer"
        assert spec.object("drawer").parent_label is None


class TestTaskNotParseableError:
    def test_freeform_instruction_raises_instead_of_fabricating_a_spec(
        self, tmp_path: Path
    ) -> None:
        profile = _profile(tmp_path)
        with pytest.raises(TaskNotParseableError):
            EpisodeSpecResolver(profile).resolve(_metadata(_UNPARSEABLE_TASK), camera_role="ext1")

    def test_meta_token_only_instruction_also_raises(self, tmp_path: Path) -> None:
        """A short phrase that is nonetheless pure meta-vocabulary must also refuse."""
        profile = _profile(tmp_path)
        metadata = _metadata("complete the suggested task")
        with pytest.raises(TaskNotParseableError):
            EpisodeSpecResolver(profile).resolve(metadata, camera_role="ext1")


class TestWithStaticSpans:
    def _resolved_spec(self, tmp_path: Path) -> EpisodeSpec:
        profile = _profile(tmp_path)
        return EpisodeSpecResolver(profile).resolve(_metadata(_REAL_TASK), camera_role="ext1")

    def test_attaches_measured_spans_by_label(self, tmp_path: Path) -> None:
        spec = self._resolved_spec(tmp_path)
        updated = spec.with_static_spans({"brick": (0, 30), "drawer": None})
        assert updated.object("brick").static_span == (0, 30)
        assert updated.object("drawer").static_span is None
        # original spec is untouched -- EpisodeSpec is frozen/immutable.
        assert spec.object("brick").static_span is None

    def test_rejects_unknown_labels(self, tmp_path: Path) -> None:
        spec = self._resolved_spec(tmp_path)
        with pytest.raises(ConfigError):
            spec.with_static_spans({"not_a_real_object": (0, 5)})

    def test_partial_update_leaves_other_labels_alone(self, tmp_path: Path) -> None:
        spec = self._resolved_spec(tmp_path)
        updated = spec.with_static_spans({"brick": (0, 12)})
        assert updated.object("brick").static_span == (0, 12)
        assert updated.object("drawer").static_span is None


class TestObjectSpecValidation:
    def test_mesh_representation_requires_a_mesh_path(self) -> None:
        with pytest.raises(ConfigError):
            ObjectSpec(
                label="brick", role=ObjectRole.MANIPULATED, prompt_candidates=("brick",),
                representation=ShapeRepresentation.MESH, mesh_path=None,
                plausible_scale_range=None, allow_scale=True,
            )

    def test_observed_surface_representation_does_not_require_a_mesh_path(self) -> None:
        spec = ObjectSpec(
            label="drawer", role=ObjectRole.FIXTURE, prompt_candidates=("drawer",),
            representation=ShapeRepresentation.OBSERVED_SURFACE, mesh_path=None,
            plausible_scale_range=None, allow_scale=False,
        )
        assert spec.mesh_path is None

    def test_empty_prompt_candidates_rejected(self) -> None:
        with pytest.raises(ConfigError):
            ObjectSpec(
                label="brick", role=ObjectRole.MANIPULATED, prompt_candidates=(),
                representation=ShapeRepresentation.OBSERVED_SURFACE, mesh_path=None,
                plausible_scale_range=None, allow_scale=True,
            )

    def test_empty_label_rejected(self) -> None:
        with pytest.raises(ConfigError):
            ObjectSpec(
                label="", role=ObjectRole.MANIPULATED, prompt_candidates=("brick",),
                representation=ShapeRepresentation.OBSERVED_SURFACE, mesh_path=None,
                plausible_scale_range=None, allow_scale=True,
            )


class TestEpisodeSpecDuplicateLabels:
    def test_duplicate_labels_rejected(self) -> None:
        obj = ObjectSpec(
            label="brick", role=ObjectRole.MANIPULATED, prompt_candidates=("brick",),
            representation=ShapeRepresentation.OBSERVED_SURFACE, mesh_path=None,
            plausible_scale_range=None, allow_scale=True,
        )
        with pytest.raises(ConfigError):
            EpisodeSpec(
                uuid="u", camera_role="ext1", camera_serial="1", scene_id="s", task="t",
                objects=(obj, obj),
            )

    def test_unknown_parent_label_rejected(self) -> None:
        obj = ObjectSpec(
            label="brick", role=ObjectRole.MANIPULATED, prompt_candidates=("brick",),
            representation=ShapeRepresentation.OBSERVED_SURFACE, mesh_path=None,
            plausible_scale_range=None, allow_scale=True, parent_label="does_not_exist",
        )
        with pytest.raises(ConfigError):
            EpisodeSpec(
                uuid="u", camera_role="ext1", camera_serial="1", scene_id="s", task="t",
                objects=(obj,),
            )