"""Integration checks against the tiny, explicitly selected source sample."""
from pathlib import Path
import pytest
from task00_recovery.source import SourceEpisode

INPUTS = Path(__file__).resolve().parents[1]/'inputs'


@pytest.mark.parametrize('episode,frame,held', [('00003990', 4000, True), ('00005200', 4299, False)])
def test_actual_grasp_state_takes_precedence_over_annotation(episode, frame, held):
    source = INPUTS/f'trajectories/task-0000/episode_{episode}.hdf5'
    meta = INPUTS/f'meta/task-0000/episodes/episode_{episode}.json'
    if not source.exists():
        pytest.skip('Selected HF test source has not been downloaded')
    s = SourceEpisode(source, meta)
    snapshot, audit = s.materialize(frame)
    robot = snapshot[0]['registry']['object_registry']['robot_r1']
    assert bool(robot['ag_obj_constraint_params']['right']) == held
    assert audit['prefix_rows'] == frame+1
    assert audit['unresolved_tails'] == {}
    assert s.n_states == s.n_actions+1
    with pytest.raises(IndexError):
        s.materialize(s.n_states)