tests/test_publish_hf_dataset.py

"""Tests for Hugging Face Dataset publishing helpers."""

from __future__ import annotations

import json
import tempfile
import unittest
from pathlib import Path

from scripts.prepare_curated_dataset import build_curated_records, write_jsonl
from scripts.publish_hf_dataset import upload_dataset, validate_dataset_file


class PublishHfDatasetTest(unittest.TestCase):
    def test_validate_dataset_file_rejects_bad_assistant_json(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            dataset_file = Path(tmp_dir) / "bad.jsonl"
            dataset_file.write_text(
                json.dumps(
                    {
                        "id": "bad",
                        "messages": [
                            {"role": "system", "content": "system"},
                            {"role": "user", "content": "user"},
                            {"role": "assistant", "content": "not json"},
                        ],
                    }
                )
                + "\n",
                encoding="utf-8",
            )

            with self.assertRaises(ValueError):
                validate_dataset_file(dataset_file)

    def test_upload_dataset_dry_run_returns_file_summary(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            dataset_file = Path(tmp_dir) / "curated_v2.jsonl"
            write_jsonl(build_curated_records(5, version="v2"), dataset_file)

            summary = upload_dataset(
                dataset_file=dataset_file,
                repo_id="qqyule/objectverse-diary-sft-curated",
                path_in_repo="objectverse_sft_curated_v2.jsonl",
                private=False,
                commit_message="Dry run",
                dry_run=True,
            )

        self.assertFalse(summary["uploaded"])
        self.assertEqual(summary["repo_id"], "qqyule/objectverse-diary-sft-curated")
        self.assertEqual(summary["path_in_repo"], "objectverse_sft_curated_v2.jsonl")
        self.assertEqual(summary["record_count"], 5)


if __name__ == "__main__":
    unittest.main()