test_dataset_hub.py

"""Tests for HuggingFace Dataset Hub browser: search, preview, download."""

from __future__ import annotations

from unittest.mock import MagicMock, patch

from typer.testing import CliRunner

from soup_cli.commands.data import app

runner = CliRunner()


# ---------------------------------------------------------------------------
# soup data search
# ---------------------------------------------------------------------------


class TestDataSearch:
    """Tests for `soup data search` subcommand."""

    @patch("soup_cli.commands.data.list_datasets")
    def test_search_basic(self, mock_list):
        """Basic search returns results table."""
        mock_ds = MagicMock()
        mock_ds.id = "teknium/OpenHermes-2.5"
        mock_ds.downloads = 50000
        mock_ds.likes = 200
        mock_ds.tags = ["en", "sft"]
        mock_list.return_value = [mock_ds]

        result = runner.invoke(app, ["search", "openhermes"])
        assert result.exit_code == 0
        assert "teknium/OpenHermes-2.5" in result.output

    @patch("soup_cli.commands.data.list_datasets")
    def test_search_no_results(self, mock_list):
        """Search with no results shows message."""
        mock_list.return_value = []

        result = runner.invoke(app, ["search", "nonexistent_dataset_xyz_123"])
        assert result.exit_code == 0
        assert "No datasets found" in result.output

    @patch("soup_cli.commands.data.list_datasets")
    def test_search_limit(self, mock_list):
        """--limit controls number of results."""
        datasets = []
        for idx in range(5):
            mock_ds = MagicMock()
            mock_ds.id = f"user/dataset-{idx}"
            mock_ds.downloads = 100
            mock_ds.likes = 10
            mock_ds.tags = []
            datasets.append(mock_ds)
        mock_list.return_value = datasets

        result = runner.invoke(app, ["search", "dataset", "--limit", "3"])
        assert result.exit_code == 0

    @patch("soup_cli.commands.data.list_datasets")
    def test_search_sort_downloads(self, mock_list):
        """--sort downloads sorts by download count."""
        mock_list.return_value = []

        result = runner.invoke(
            app, ["search", "code", "--sort", "downloads"]
        )
        assert result.exit_code == 0
        mock_list.assert_called_once()
        call_kwargs = mock_list.call_args
        assert call_kwargs[1].get("sort") == "downloads"

    @patch("soup_cli.commands.data.list_datasets")
    def test_search_sort_likes(self, mock_list):
        """--sort likes sorts by likes."""
        mock_list.return_value = []

        result = runner.invoke(app, ["search", "code", "--sort", "likes"])
        assert result.exit_code == 0
        call_kwargs = mock_list.call_args
        assert call_kwargs[1].get("sort") == "likes"

    @patch("soup_cli.commands.data.list_datasets")
    def test_search_sort_invalid(self, mock_list):
        """Invalid --sort value is rejected."""
        result = runner.invoke(app, ["search", "code", "--sort", "invalid"])
        assert result.exit_code != 0

    @patch("soup_cli.commands.data.list_datasets")
    def test_search_huggingface_hub_not_installed(self, mock_list):
        """Graceful error when huggingface_hub is missing."""
        mock_list.side_effect = ImportError("No module named 'huggingface_hub'")

        result = runner.invoke(app, ["search", "code"])
        assert result.exit_code == 1
        assert "huggingface" in result.output.lower()


# ---------------------------------------------------------------------------
# soup data preview
# ---------------------------------------------------------------------------


class TestDataPreview:
    """Tests for `soup data preview` subcommand."""

    @patch("soup_cli.commands.data._hf_dataset_info")
    def test_preview_basic(self, mock_info):
        """Preview shows dataset info table."""
        mock_info.return_value = {
            "id": "teknium/OpenHermes-2.5",
            "description": "A large collection of instruction pairs.",
            "downloads": 50000,
            "likes": 200,
            "size_bytes": 2_000_000_000,
            "splits": {"train": 1_000_000, "test": 10_000},
            "features": ["conversations"],
            "tags": ["en", "sft"],
        }

        result = runner.invoke(app, ["preview", "teknium/OpenHermes-2.5"])
        assert result.exit_code == 0
        assert "teknium/OpenHermes-2.5" in result.output

    @patch("soup_cli.commands.data._hf_dataset_info")
    def test_preview_not_found(self, mock_info):
        """Preview of nonexistent dataset shows error."""
        mock_info.side_effect = ValueError("Dataset not found")

        result = runner.invoke(app, ["preview", "nonexistent/dataset"])
        assert result.exit_code == 1
        assert "not found" in result.output.lower() or "error" in result.output.lower()

    @patch("soup_cli.commands.data._hf_dataset_info")
    def test_preview_shows_splits(self, mock_info):
        """Preview shows split information."""
        mock_info.return_value = {
            "id": "test/ds",
            "description": "Test",
            "downloads": 100,
            "likes": 5,
            "size_bytes": 1000,
            "splits": {"train": 500, "validation": 100},
            "features": ["text"],
            "tags": [],
        }

        result = runner.invoke(app, ["preview", "test/ds"])
        assert result.exit_code == 0
        assert "train" in result.output

    @patch("soup_cli.commands.data._hf_dataset_info")
    def test_preview_shows_features(self, mock_info):
        """Preview shows feature columns."""
        mock_info.return_value = {
            "id": "test/ds",
            "description": "Test",
            "downloads": 100,
            "likes": 5,
            "size_bytes": 1000,
            "splits": {"train": 500},
            "features": ["instruction", "output", "input"],
            "tags": [],
        }

        result = runner.invoke(app, ["preview", "test/ds"])
        assert result.exit_code == 0


# ---------------------------------------------------------------------------
# soup data download
# ---------------------------------------------------------------------------


class TestDataDownload:
    """Tests for `soup data download` subcommand."""

    @patch("soup_cli.commands.data._hf_download_dataset")
    def test_download_basic(self, mock_download, tmp_path, monkeypatch):
        """Basic download writes JSONL output."""
        monkeypatch.chdir(tmp_path)
        output_file = tmp_path / "data.jsonl"
        mock_download.return_value = [
            {"instruction": "What is 2+2?", "output": "4"},
            {"instruction": "Hello", "output": "Hi there"},
        ]

        result = runner.invoke(
            app, ["download", "test/dataset", "-o", str(output_file)]
        )
        assert result.exit_code == 0
        assert output_file.exists()
        lines = output_file.read_text(encoding="utf-8").strip().split("\n")
        assert len(lines) == 2

    @patch("soup_cli.commands.data._hf_download_dataset")
    def test_download_with_split(self, mock_download, tmp_path, monkeypatch):
        """--split flag is passed to download function."""
        monkeypatch.chdir(tmp_path)
        output_file = tmp_path / "data.jsonl"
        mock_download.return_value = [{"text": "hello"}]

        result = runner.invoke(
            app,
            ["download", "test/dataset", "--split", "train[:100]", "-o", str(output_file)],
        )
        assert result.exit_code == 0
        mock_download.assert_called_once()
        call_kwargs = mock_download.call_args
        assert call_kwargs[1].get("split") == "train[:100]"

    @patch("soup_cli.commands.data._hf_download_dataset")
    def test_download_with_format_conversion(self, mock_download, tmp_path, monkeypatch):
        """--format converts downloaded data."""
        monkeypatch.chdir(tmp_path)
        output_file = tmp_path / "data.jsonl"
        # Alpaca format input
        mock_download.return_value = [
            {"instruction": "Q1", "input": "", "output": "A1"},
        ]

        result = runner.invoke(
            app,
            [
                "download", "test/dataset",
                "--format", "sharegpt",
                "-o", str(output_file),
            ],
        )
        assert result.exit_code == 0
        assert output_file.exists()

    @patch("soup_cli.commands.data._hf_download_dataset")
    def test_download_empty_dataset(self, mock_download, tmp_path, monkeypatch):
        """Empty download result shows error."""
        monkeypatch.chdir(tmp_path)
        mock_download.return_value = []
        output_file = tmp_path / "data.jsonl"

        result = runner.invoke(
            app, ["download", "test/empty", "-o", str(output_file)]
        )
        assert result.exit_code == 1
        assert "empty" in result.output.lower() or "no data" in result.output.lower()

    @patch("soup_cli.commands.data._hf_download_dataset")
    def test_download_error_handling(self, mock_download, tmp_path):
        """Download failure shows friendly error."""
        mock_download.side_effect = ValueError("Dataset not found on HuggingFace Hub")
        output_file = tmp_path / "data.jsonl"

        result = runner.invoke(
            app, ["download", "nonexistent/dataset", "-o", str(output_file)]
        )
        assert result.exit_code == 1

    @patch("soup_cli.commands.data._hf_download_dataset")
    def test_download_output_path_traversal(self, mock_download, tmp_path):
        """Output path must stay under cwd."""
        mock_download.return_value = [{"text": "hello"}]

        result = runner.invoke(
            app,
            ["download", "test/dataset", "-o", "/etc/passwd"],
        )
        assert result.exit_code == 1
        assert "current working directory" in result.output.lower()

    @patch("soup_cli.commands.data._hf_download_dataset")
    def test_download_default_output(self, mock_download, tmp_path, monkeypatch):
        """Default output file uses dataset name."""
        monkeypatch.chdir(tmp_path)
        mock_download.return_value = [{"text": "hello"}]

        result = runner.invoke(
            app, ["download", "user/my-dataset"]
        )
        assert result.exit_code == 0
        # Default should be my-dataset.jsonl in cwd
        expected = tmp_path / "my-dataset.jsonl"
        assert expected.exists()

    @patch("soup_cli.commands.data._hf_download_dataset")
    def test_download_samples_limit(self, mock_download, tmp_path, monkeypatch):
        """--samples limits number of downloaded rows."""
        monkeypatch.chdir(tmp_path)
        output_file = tmp_path / "data.jsonl"
        mock_download.return_value = [
            {"text": f"row {idx}"} for idx in range(10)
        ]

        result = runner.invoke(
            app, ["download", "test/dataset", "--samples", "5", "-o", str(output_file)]
        )
        assert result.exit_code == 0
        lines = output_file.read_text(encoding="utf-8").strip().split("\n")
        assert len(lines) == 5


# ---------------------------------------------------------------------------
# Helper function tests
# ---------------------------------------------------------------------------


class TestHelperFunctions:
    """Tests for internal helper functions."""

    def test_format_size_bytes(self):
        """_format_size_bytes formats various sizes correctly."""
        from soup_cli.commands.data import _format_size_bytes

        assert _format_size_bytes(0) == "0 B"
        assert _format_size_bytes(500) == "500 B"
        assert "KB" in _format_size_bytes(1024)
        assert "MB" in _format_size_bytes(1024 * 1024)
        assert "GB" in _format_size_bytes(1024 * 1024 * 1024)

    def test_format_size_bytes_none(self):
        """_format_size_bytes handles None."""
        from soup_cli.commands.data import _format_size_bytes

        assert _format_size_bytes(None) == "unknown"

    def test_format_count(self):
        """_format_count formats large numbers."""
        from soup_cli.commands.data import _format_count

        assert _format_count(500) == "500"
        assert "K" in _format_count(1500)
        assert "M" in _format_count(1_500_000)


# ---------------------------------------------------------------------------
# Security edge cases
# ---------------------------------------------------------------------------


class TestSecurityEdgeCases:
    """Security-relevant edge case tests."""

    def test_download_samples_over_limit(self):
        """--samples above 1M is rejected."""
        result = runner.invoke(
            app,
            ["download", "test/dataset", "--samples", "2000000", "-o", "out.jsonl"],
        )
        assert result.exit_code == 1
        assert "1,000,000" in result.output

    @patch("soup_cli.commands.data._hf_dataset_info")
    def test_preview_huggingface_hub_not_installed(self, mock_info):
        """Preview gracefully handles missing huggingface_hub."""
        mock_info.side_effect = ImportError("No module named 'huggingface_hub'")
        result = runner.invoke(app, ["preview", "test/ds"])
        assert result.exit_code == 1
        assert "huggingface" in result.output.lower()