test_data_augment.py

"""Tests for data augmentation (Part F of v0.25.0)."""

import json

import pytest
from typer.testing import CliRunner

from soup_cli.cli import app

runner = CliRunner()


# ---------------------------------------------------------------------------
# Strategy unit tests (with a mocked provider)
# ---------------------------------------------------------------------------

class FakeProvider:
    """Deterministic fake LLM provider for testing."""

    def __init__(self):
        self.calls = 0

    def generate(self, prompt: str, max_tokens: int = 512) -> str:
        self.calls += 1
        return f"rewritten-{self.calls}"


class TestAugmentStrategies:
    def test_rephrase_basic(self):
        from soup_cli.data.augment import augment_rephrase

        examples = [
            {"instruction": "What is Python?", "output": "Python is a language."},
        ]
        provider = FakeProvider()
        augmented = augment_rephrase(examples, provider=provider, count=2)
        # 1 original × 2 count × 2 string fields (instruction + output) = 4 provider calls
        assert provider.calls == 4
        # 1 original → 2 new augmentations
        assert len(augmented) == 2
        for row in augmented:
            assert "instruction" in row or "output" in row

    def test_translate_produces_langs(self):
        from soup_cli.data.augment import augment_translate

        examples = [
            {"instruction": "Hello", "output": "Hi"},
        ]
        provider = FakeProvider()
        augmented = augment_translate(
            examples, provider=provider, languages=["ru", "zh"]
        )
        # 1 example × 2 languages × 2 string fields = 4 provider calls
        assert provider.calls == 4
        assert len(augmented) == 2

    def test_style_produces_styles(self):
        from soup_cli.data.augment import augment_style

        examples = [
            {
                "instruction": "Explain recursion.",
                "output": "It's when a function calls itself.",
            },
        ]
        provider = FakeProvider()
        augmented = augment_style(
            examples, provider=provider, styles=["formal", "casual"]
        )
        # 1 example × 2 styles × 2 string fields = 4 provider calls
        assert provider.calls == 4
        assert len(augmented) == 2

    def test_count_cap_enforced(self):
        """count > 10 is rejected at strategy level (not just CLI)."""
        from soup_cli.data.augment import augment_rephrase

        examples = [{"instruction": "x", "output": "y"}]
        provider = FakeProvider()
        with pytest.raises(ValueError):
            augment_rephrase(examples, provider=provider, count=11)

    def test_empty_examples_returns_empty(self):
        from soup_cli.data.augment import (
            augment_rephrase,
            augment_style,
            augment_translate,
        )

        provider = FakeProvider()
        assert augment_rephrase([], provider=provider, count=2) == []
        assert augment_translate([], provider=provider, languages=["ru"]) == []
        assert augment_style([], provider=provider, styles=["formal"]) == []
        assert provider.calls == 0

    def test_translate_empty_languages_rejected(self):
        from soup_cli.data.augment import augment_translate

        with pytest.raises(ValueError):
            augment_translate([{"x": "y"}], provider=FakeProvider(), languages=[])

    def test_style_empty_styles_rejected(self):
        from soup_cli.data.augment import augment_style

        with pytest.raises(ValueError):
            augment_style([{"x": "y"}], provider=FakeProvider(), styles=[])

    def test_provider_exception_propagates(self):
        """When provider raises, augment propagates the error (fail-loud)."""
        from soup_cli.data.augment import augment_rephrase

        class FailingProvider:
            def generate(self, prompt: str, max_tokens: int = 512) -> str:
                raise RuntimeError("provider down")

        with pytest.raises(RuntimeError, match="provider down"):
            augment_rephrase(
                [{"instruction": "x"}], provider=FailingProvider(), count=1,
            )


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

class TestAugmentCLI:
    def _write_data(self, tmp_path):
        rows = [
            {"instruction": f"q{i}", "output": f"a{i}"} for i in range(3)
        ]
        path = tmp_path / "data.jsonl"
        path.write_text(
            "\n".join(json.dumps(r) for r in rows) + "\n", encoding="utf-8"
        )
        return path

    def test_augment_help(self):
        result = runner.invoke(app, ["data", "augment", "--help"])
        assert result.exit_code == 0
        assert "augment" in result.output.lower()

    def test_augment_missing_file(self, tmp_path, monkeypatch):
        monkeypatch.chdir(tmp_path)
        result = runner.invoke(app, [
            "data", "augment",
            "--input", "nonexistent.jsonl",
            "--output", "out.jsonl",
            "--strategy", "rephrase",
            "--provider", "ollama",
        ])
        assert result.exit_code != 0

    def test_augment_count_cap(self, tmp_path, monkeypatch):
        """CLI rejects count > 10."""
        monkeypatch.chdir(tmp_path)
        path = self._write_data(tmp_path)
        result = runner.invoke(app, [
            "data", "augment",
            "--input", str(path.name),
            "--output", "out.jsonl",
            "--strategy", "rephrase",
            "--count", "15",
            "--provider", "ollama",
        ])
        assert result.exit_code != 0

    def test_input_path_traversal_blocked(self, tmp_path, monkeypatch):
        monkeypatch.chdir(tmp_path)
        result = runner.invoke(app, [
            "data", "augment",
            "--input", "../../../etc/passwd",
            "--output", "out.jsonl",
            "--strategy", "rephrase",
            "--provider", "ollama",
        ])
        assert result.exit_code != 0

    def test_output_path_traversal_blocked(self, tmp_path, monkeypatch):
        monkeypatch.chdir(tmp_path)
        path = self._write_data(tmp_path)
        result = runner.invoke(app, [
            "data", "augment",
            "--input", str(path.name),
            "--output", "../../evil.jsonl",
            "--strategy", "rephrase",
            "--provider", "ollama",
        ])
        assert result.exit_code != 0


# ---------------------------------------------------------------------------
# Strategy validation
# ---------------------------------------------------------------------------

class TestStrategyValidation:
    def test_unknown_strategy(self, tmp_path, monkeypatch):
        monkeypatch.chdir(tmp_path)
        path = tmp_path / "data.jsonl"
        path.write_text('{"instruction": "x", "output": "y"}\n', encoding="utf-8")
        result = runner.invoke(app, [
            "data", "augment",
            "--input", str(path.name),
            "--output", "out.jsonl",
            "--strategy", "bogus",
            "--provider", "ollama",
        ])
        assert result.exit_code != 0


if __name__ == "__main__":
    pytest.main([__file__, "-v"])