test_data_sample.py

"""Tests for soup data sample — intelligent dataset sampling."""

import json
from pathlib import Path

from typer.testing import CliRunner

from soup_cli.cli import app

runner = CliRunner()


def _create_jsonl(tmp_path: Path, filename: str, num_rows: int) -> Path:
    """Helper to create a JSONL file with sample data."""
    file_path = tmp_path / filename
    with open(file_path, "w", encoding="utf-8") as fh:
        for idx in range(num_rows):
            row = {
                "instruction": f"Question {idx}",
                "input": f"Context for question {idx}" * (idx + 1),
                "output": f"Answer {idx}" * (idx + 1),
            }
            fh.write(json.dumps(row) + "\n")
    return file_path


class TestSampleCLI:
    """Test soup data sample CLI command."""

    def test_random_sample_n(self, tmp_path, monkeypatch):
        """Sample N rows with random strategy."""
        monkeypatch.chdir(tmp_path)
        input_path = _create_jsonl(tmp_path, "big.jsonl", 100)
        output_path = tmp_path / "small.jsonl"
        result = runner.invoke(app, [
            "data", "sample",
            str(input_path),
            "--output", str(output_path),
            "--n", "10",
        ])
        assert result.exit_code == 0
        with open(output_path, encoding="utf-8") as fh:
            rows = [json.loads(line) for line in fh]
        assert len(rows) == 10

    def test_random_sample_pct(self, tmp_path, monkeypatch):
        """Sample by percentage."""
        monkeypatch.chdir(tmp_path)
        input_path = _create_jsonl(tmp_path, "big.jsonl", 100)
        output_path = tmp_path / "small.jsonl"
        result = runner.invoke(app, [
            "data", "sample",
            str(input_path),
            "--output", str(output_path),
            "--pct", "10",
        ])
        assert result.exit_code == 0
        with open(output_path, encoding="utf-8") as fh:
            rows = [json.loads(line) for line in fh]
        assert len(rows) == 10

    def test_sample_n_larger_than_dataset(self, tmp_path, monkeypatch):
        """When n > dataset size, return all rows."""
        monkeypatch.chdir(tmp_path)
        input_path = _create_jsonl(tmp_path, "small.jsonl", 5)
        output_path = tmp_path / "out.jsonl"
        result = runner.invoke(app, [
            "data", "sample",
            str(input_path),
            "--output", str(output_path),
            "--n", "100",
        ])
        assert result.exit_code == 0
        with open(output_path, encoding="utf-8") as fh:
            rows = [json.loads(line) for line in fh]
        assert len(rows) == 5

    def test_sample_with_seed(self, tmp_path, monkeypatch):
        """Seed produces deterministic output."""
        monkeypatch.chdir(tmp_path)
        input_path = _create_jsonl(tmp_path, "data.jsonl", 50)
        out1 = tmp_path / "out1.jsonl"
        out2 = tmp_path / "out2.jsonl"

        runner.invoke(app, [
            "data", "sample", str(input_path),
            "--output", str(out1), "--n", "10", "--seed", "42",
        ])
        runner.invoke(app, [
            "data", "sample", str(input_path),
            "--output", str(out2), "--n", "10", "--seed", "42",
        ])

        with open(out1, encoding="utf-8") as fh:
            rows1 = fh.readlines()
        with open(out2, encoding="utf-8") as fh:
            rows2 = fh.readlines()
        assert rows1 == rows2

    def test_diverse_strategy(self, tmp_path, monkeypatch):
        """Diverse strategy returns requested number of samples."""
        monkeypatch.chdir(tmp_path)
        input_path = _create_jsonl(tmp_path, "data.jsonl", 50)
        output_path = tmp_path / "diverse.jsonl"
        result = runner.invoke(app, [
            "data", "sample", str(input_path),
            "--output", str(output_path),
            "--n", "10",
            "--strategy", "diverse",
        ])
        assert result.exit_code == 0
        with open(output_path, encoding="utf-8") as fh:
            rows = [json.loads(line) for line in fh]
        assert len(rows) == 10

    def test_hard_strategy(self, tmp_path, monkeypatch):
        """Hard strategy returns requested number of samples."""
        monkeypatch.chdir(tmp_path)
        input_path = _create_jsonl(tmp_path, "data.jsonl", 50)
        output_path = tmp_path / "hard.jsonl"
        result = runner.invoke(app, [
            "data", "sample", str(input_path),
            "--output", str(output_path),
            "--n", "10",
            "--strategy", "hard",
        ])
        assert result.exit_code == 0
        with open(output_path, encoding="utf-8") as fh:
            rows = [json.loads(line) for line in fh]
        assert len(rows) == 10

    def test_missing_input_file(self):
        """Should fail for nonexistent input file."""
        result = runner.invoke(app, [
            "data", "sample", "nonexistent.jsonl",
            "--output", "out.jsonl", "--n", "10",
        ])
        assert result.exit_code != 0

    def test_empty_input(self, tmp_path):
        """Should handle empty input file gracefully."""
        empty_file = tmp_path / "empty.jsonl"
        empty_file.write_text("")
        output_path = tmp_path / "out.jsonl"
        result = runner.invoke(app, [
            "data", "sample", str(empty_file),
            "--output", str(output_path), "--n", "10",
        ])
        assert result.exit_code != 0

    def test_default_output_name(self, tmp_path):
        """v0.40.1 — default filename embeds the strategy to prevent
        overwrite when running successive `random`/`diverse`/`hard` passes.
        """
        input_path = _create_jsonl(tmp_path, "data.jsonl", 20)
        result = runner.invoke(app, [
            "data", "sample", str(input_path), "--n", "5",
        ])
        assert result.exit_code == 0
        expected_output = tmp_path / "data_sampled_random.jsonl"
        assert expected_output.exists()
        with open(expected_output, encoding="utf-8") as fh:
            rows = [json.loads(line) for line in fh]
        assert len(rows) == 5

    def test_must_specify_n_or_pct(self, tmp_path):
        """Should fail if neither --n nor --pct specified."""
        input_path = _create_jsonl(tmp_path, "data.jsonl", 20)
        result = runner.invoke(app, [
            "data", "sample", str(input_path),
        ])
        assert result.exit_code != 0

    def test_invalid_strategy(self, tmp_path):
        """Unknown strategy shows error."""
        input_path = _create_jsonl(tmp_path, "data.jsonl", 20)
        result = runner.invoke(app, [
            "data", "sample", str(input_path),
            "--n", "5", "--strategy", "nonexistent",
        ])
        assert result.exit_code != 0


class TestSampleSecurity:
    """Security tests for sample command."""

    def test_sample_output_path_traversal(self, tmp_path):
        """--output with path traversal should be rejected."""
        input_path = _create_jsonl(tmp_path, "data.jsonl", 20)
        result = runner.invoke(app, [
            "data", "sample", str(input_path),
            "--output", "../../evil.jsonl", "--n", "5",
        ])
        assert result.exit_code != 0


class TestSampleStrategies:
    """Test sampling strategy functions directly."""

    def test_random_sample(self):
        from soup_cli.commands.data import _sample_random

        data = [{"text": f"row {idx}"} for idx in range(100)]
        result = _sample_random(data, 10, seed=42)
        assert len(result) == 10
        # All sampled items should be from original data
        for item in result:
            assert item in data

    def test_diverse_sample(self):
        from soup_cli.commands.data import _sample_diverse

        data = [
            {"text": "Python is a programming language" * (idx + 1)}
            for idx in range(50)
        ]
        result = _sample_diverse(data, 10, seed=42)
        assert len(result) == 10

    def test_hard_sample(self):
        from soup_cli.commands.data import _sample_hard

        # Create data with varying lengths (proxy for difficulty)
        data = [
            {"text": "word " * (idx + 1)}
            for idx in range(50)
        ]
        result = _sample_hard(data, 10)
        assert len(result) == 10
        # Hard samples should be longer (more difficult)
        avg_len = sum(len(str(row)) for row in result) / len(result)
        all_avg = sum(len(str(row)) for row in data) / len(data)
        assert avg_len >= all_avg  # hard samples are above average length