tests/test_publish.py

"""Tests for fpgm.datagen.publish.DatasetPublisher."""

from __future__ import annotations

import json
from pathlib import Path

import h5py
import numpy as np
import pytest

from fpgm.config_datagen import DatagenProfile
from fpgm.datagen.publish import DatasetPublisher

REPO_ROOT = Path(__file__).resolve().parents[1]


def _write_profile_yaml(tmp_path: Path, *, include_debug_videos: bool = True) -> Path:
    import yaml

    raw = {
        "paths": {
            "droid_raw_root": str(tmp_path / "droid_raw"),
            "flows_root": str(tmp_path / "flows"),
            "cameras_root": str(tmp_path / "cameras"),
            "output_root": str(tmp_path / "outputs"),
            "checkpoints_root": str(tmp_path / "checkpoints"),
            "urdf": str(tmp_path / "fake.urdf"),
            "python_executable": str(tmp_path / "python"),
        },
        "publish": {
            "bucket_url": "hf://buckets/test-org/test-bucket",
            "master_includes": ["poses.npz", "events.json", "*_masks.h5"],
            "vace_includes": ["target.mp4", "sample.json"],
            "include_debug_videos": include_debug_videos,
        },
    }
    yaml_path = tmp_path / "profile.yaml"
    yaml_path.write_text(yaml.safe_dump(raw))
    return yaml_path


def _build_episode(profile: DatagenProfile, uuid: str, camera_serial: str = "cam1") -> None:
    """A minimal but realistic-shaped episode/camera output tree."""
    master_dir = profile.paths.master_dir(uuid, camera_serial)
    master_dir.mkdir(parents=True)
    np.savez(
        master_dir / "poses.npz",
        labels=np.array(["brick"]),
        brick__pose_source=np.array([0, 0, 1, 3], dtype=np.int64),
    )
    (master_dir / "events.json").write_text("{}")
    (master_dir / "object_0_masks.h5").write_bytes(b"masksdata")
    (master_dir / "object_poses_debug.mp4").write_bytes(b"debugvideo")
    # robot_buffers -- large intermediates, must never be staged
    buf = master_dir / "robot_buffers"
    buf.mkdir()
    (buf / "depth_mm.npy").write_bytes(b"x" * 10_000)
    (buf / "seg.npy").write_bytes(b"y" * 5_000)
    (buf / "meta.json").write_text("{}")

    camera_dir = profile.paths.camera_dir(uuid, camera_serial)
    window_dir = camera_dir / "vace" / "window_00000_00010"
    window_dir.mkdir(parents=True)
    (window_dir / "target.mp4").write_bytes(b"targetvideo")
    sample = {
        "hole_fraction_mean": 0.02, "hole_fraction_max": 0.05,
        "overlaps_pose_gap": False, "gates_passed": {"resolution_contract_hw16": True},
    }
    (window_dir / "sample.json").write_text(json.dumps(sample))
    (window_dir / "not_included.txt").write_text("should not be staged")

    # flows-h5 metadata, so the manifest can fill scene_id/task/camera_role
    flows_path = profile.paths.flows_h5(uuid)
    flows_path.parent.mkdir(parents=True, exist_ok=True)
    with h5py.File(flows_path, "w") as fh:
        fh.attrs["uuid"] = uuid
        fh.attrs["scene_id"] = "8756300955"
        fh.attrs["current_task"] = "Put brick in drawer shelf and close drawer"
        fh.attrs["ext1_cam_serial"] = camera_serial


class TestDiscovery:
    def test_finds_episode_camera_pairs_with_master_dir(self, tmp_path):
        yaml_path = _write_profile_yaml(tmp_path)
        profile = DatagenProfile.from_yaml(yaml_path)
        _build_episode(profile, "ep1", "cam1")
        _build_episode(profile, "ep2", "cam2")
        # a directory with no master/ must not be picked up
        (profile.paths.output_root / "ep3" / "not_a_camera").mkdir(parents=True)

        publisher = DatasetPublisher(profile)
        pairs = publisher.discover_episode_cameras()
        assert sorted(pairs) == [("ep1", "cam1"), ("ep2", "cam2")]

    def test_filters_by_episode_uuids(self, tmp_path):
        yaml_path = _write_profile_yaml(tmp_path)
        profile = DatagenProfile.from_yaml(yaml_path)
        _build_episode(profile, "ep1", "cam1")
        _build_episode(profile, "ep2", "cam2")

        publisher = DatasetPublisher(profile)
        pairs = publisher.discover_episode_cameras(episode_uuids=["ep2"])
        assert pairs == [("ep2", "cam2")]


class TestStaging:
    def test_stages_only_included_files_excludes_robot_buffers_npy(self, tmp_path):
        yaml_path = _write_profile_yaml(tmp_path, include_debug_videos=False)
        profile = DatagenProfile.from_yaml(yaml_path)
        _build_episode(profile, "ep1", "cam1")

        publisher = DatasetPublisher(profile)
        report = publisher.stage(tmp_path / "staging")

        assert report.n_episodes == 1
        staged_files = {
            str(p.relative_to(report.staging_dir))
            for p in report.staging_dir.rglob("*") if p.is_file()
        }
        assert "ep1/cam1/master/poses.npz" in staged_files
        assert "ep1/cam1/master/events.json" in staged_files
        assert "ep1/cam1/master/object_0_masks.h5" in staged_files
        assert "ep1/cam1/vace/window_00000_00010/target.mp4" in staged_files
        assert "ep1/cam1/vace/window_00000_00010/sample.json" in staged_files
        assert "dataset_manifest.json" in staged_files
        assert "README.md" in staged_files

        # deliberately excluded
        assert "ep1/cam1/master/robot_buffers/depth_mm.npy" not in staged_files
        assert "ep1/cam1/master/robot_buffers/seg.npy" not in staged_files
        assert "ep1/cam1/vace/window_00000_00010/not_included.txt" not in staged_files
        # debug videos off in this profile
        assert "ep1/cam1/master/object_poses_debug.mp4" not in staged_files

    def test_include_debug_videos_toggle(self, tmp_path):
        yaml_path = _write_profile_yaml(tmp_path, include_debug_videos=True)
        profile = DatagenProfile.from_yaml(yaml_path)
        _build_episode(profile, "ep1", "cam1")

        publisher = DatasetPublisher(profile)
        report = publisher.stage(tmp_path / "staging")
        staged_files = {
            str(p.relative_to(report.staging_dir))
            for p in report.staging_dir.rglob("*") if p.is_file()
        }
        assert "ep1/cam1/master/object_poses_debug.mp4" in staged_files

    def test_manifest_has_uuid_camera_scene_task_gates_windows_and_file_hashes(self, tmp_path):
        yaml_path = _write_profile_yaml(tmp_path)
        profile = DatagenProfile.from_yaml(yaml_path)
        _build_episode(profile, "ep1", "cam1")

        publisher = DatasetPublisher(profile)
        report = publisher.stage(tmp_path / "staging")
        manifest = json.loads(report.manifest_path.read_text())

        assert manifest["n_episodes"] == 1
        entry = manifest["episodes"][0]
        assert entry["uuid"] == "ep1"
        assert entry["camera_serial"] == "cam1"
        assert entry["camera_role"] == "ext1"
        assert entry["scene_id"] == "8756300955"
        assert entry["task"] == "Put brick in drawer shelf and close drawer"
        assert isinstance(entry["gates"], dict)
        assert len(entry["windows"]) == 1
        assert entry["windows"][0]["window"] == "window_00000_00010"

        files_by_path = {f["path"]: f for f in entry["files"]}
        poses_entry = files_by_path["master/poses.npz"]
        raw = (profile.paths.master_dir("ep1", "cam1") / "poses.npz").read_bytes()
        import hashlib
        assert poses_entry["bytes"] == len(raw)
        assert poses_entry["sha256"] == hashlib.sha256(raw).hexdigest()
        assert entry["gates"]["pose_provenance"]["brick"] == {
            "PNP": 2, "FK_ATTACH": 1, "INTERP": 0, "GAP": 1,
        }

    def test_byte_and_file_totals_match_staged_tree(self, tmp_path):
        yaml_path = _write_profile_yaml(tmp_path, include_debug_videos=False)
        profile = DatagenProfile.from_yaml(yaml_path)
        _build_episode(profile, "ep1", "cam1")
        _build_episode(profile, "ep2", "cam2")

        publisher = DatasetPublisher(profile)
        report = publisher.stage(tmp_path / "staging")

        on_disk_files = [p for p in report.staging_dir.rglob("*") if p.is_file()]
        assert report.n_files == len(on_disk_files)
        assert report.total_bytes == sum(p.stat().st_size for p in on_disk_files)

    def test_restaging_replaces_previous_contents(self, tmp_path):
        yaml_path = _write_profile_yaml(tmp_path)
        profile = DatagenProfile.from_yaml(yaml_path)
        _build_episode(profile, "ep1", "cam1")
        publisher = DatasetPublisher(profile)
        staging_dir = tmp_path / "staging"
        staging_dir.mkdir()
        (staging_dir / "stale_leftover.txt").write_text("old")

        report = publisher.stage(staging_dir)
        assert not (staging_dir / "stale_leftover.txt").exists()
        assert report.n_episodes == 1


class _FakeSyncPlan:
    def __init__(self, **summary_values):
        self._summary_values = summary_values

    def summary(self) -> dict:
        return self._summary_values


class _FakeHfApi:
    def __init__(self, *, token=None):
        self.init_token = token
        self.sync_calls: list[dict] = []

    def sync_bucket(self, **kwargs):
        self.sync_calls.append(kwargs)
        return _FakeSyncPlan(uploads=3, downloads=0, deletes=0, skips=0, total_size=12345)


class TestPublish:
    def test_dry_run_never_uploads_and_does_not_use_cached_token(self, tmp_path, monkeypatch):
        monkeypatch.delenv("HF_TOKEN", raising=False)
        yaml_path = _write_profile_yaml(tmp_path)
        profile = DatagenProfile.from_yaml(yaml_path)
        _build_episode(profile, "ep1", "cam1")
        publisher = DatasetPublisher(profile)
        report = publisher.stage(tmp_path / "staging")

        fake_api = _FakeHfApi()
        publisher_with_fake = DatasetPublisher(profile, api=fake_api)
        plan = publisher_with_fake.publish(report.staging_dir, dry_run=True)

        assert plan.summary()["uploads"] == 3
        assert len(fake_api.sync_calls) == 1
        call = fake_api.sync_calls[0]
        assert call["dry_run"] is True
        assert call["dest"] == "hf://buckets/test-org/test-bucket"
        assert call["source"] == str(report.staging_dir)
        # HF_TOKEN unset -> token=False passed explicitly, never omitted
        # (omitting it is how huggingface_hub falls back to a cached token).
        assert call["token"] is False

    def test_hf_token_from_environment_is_passed_explicitly(self, tmp_path, monkeypatch):
        monkeypatch.setenv("HF_TOKEN", "test-token-value")
        yaml_path = _write_profile_yaml(tmp_path)
        profile = DatagenProfile.from_yaml(yaml_path)
        _build_episode(profile, "ep1", "cam1")
        publisher = DatasetPublisher(profile)
        report = publisher.stage(tmp_path / "staging")

        fake_api = _FakeHfApi()
        publisher_with_fake = DatasetPublisher(profile, api=fake_api)
        publisher_with_fake.publish(report.staging_dir, dry_run=True)

        assert fake_api.sync_calls[0]["token"] == "test-token-value"

    def test_real_upload_without_token_raises(self, tmp_path, monkeypatch):
        monkeypatch.delenv("HF_TOKEN", raising=False)
        yaml_path = _write_profile_yaml(tmp_path)
        profile = DatagenProfile.from_yaml(yaml_path)
        _build_episode(profile, "ep1", "cam1")
        publisher = DatasetPublisher(profile)
        report = publisher.stage(tmp_path / "staging")

        fake_api = _FakeHfApi()
        publisher_with_fake = DatasetPublisher(profile, api=fake_api)
        with pytest.raises(RuntimeError, match="HF_TOKEN"):
            publisher_with_fake.publish(report.staging_dir, dry_run=False)
        assert fake_api.sync_calls == []  # never even attempted