soup/tests/test_server_generate.py

"""Tests for local server data generation provider (--provider server)."""

from unittest.mock import MagicMock
from unittest.mock import patch as mock_patch

import pytest

# ─── Provider Validation Tests ───────────────────────────────────────────


class TestServerProviderValidation:
    """Test that 'server' is a valid provider for soup data generate."""

    def test_server_provider_routes_without_error(self):
        """'server' should route to _generate_server without error."""
        from soup_cli.commands.generate import _generate_batch

        with mock_patch(
            "soup_cli.commands.generate._generate_server",
            return_value=[],
        ):
            result = _generate_batch(
                prompt="test", count=1, fmt="alpaca",
                provider="server", model_name="m", api_key=None,
                api_base=None, temperature=0.8, seed_examples=[],
            )
        assert result == []

    def test_invalid_provider_rejected(self):
        """Invalid providers should cause an exit."""
        from typer.testing import CliRunner

        from soup_cli.cli import app

        runner = CliRunner()
        result = runner.invoke(app, [
            "data", "generate",
            "--prompt", "test",
            "--provider", "invalid_provider",
            "--count", "1",
        ])
        assert result.exit_code != 0

    def test_server_provider_in_help_text(self):
        """The help text should mention 'server' provider."""
        from typer.testing import CliRunner

        from soup_cli.cli import app

        runner = CliRunner()
        result = runner.invoke(app, ["data", "generate", "--help"])
        assert "server" in result.output


# ─── _generate_server Function Tests ─────────────────────────────────────


class TestGenerateServer:
    """Test the _generate_server function."""

    def test_generate_server_calls_httpx(self):
        """_generate_server should call httpx.post with correct URL."""
        from soup_cli.commands.generate import _generate_server

        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = {
            "choices": [
                {"message": {"content": '[{"instruction": "test", "output": "response"}]'}}
            ]
        }

        with mock_patch("httpx.post", return_value=mock_response) as mock_post:
            _generate_server(
                prompt="Generate test data",
                count=1,
                fmt="alpaca",
                model_name="test-model",
                api_base=None,
                temperature=0.8,
                seed_examples=[],
            )

        mock_post.assert_called_once()
        call_url = mock_post.call_args[0][0]
        assert "localhost:8000" in call_url
        assert "/chat/completions" in call_url

    def test_generate_server_no_auth_header(self):
        """_generate_server should NOT include Authorization header."""
        from soup_cli.commands.generate import _generate_server

        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = {
            "choices": [
                {"message": {"content": '[]'}}
            ]
        }

        with mock_patch("httpx.post", return_value=mock_response) as mock_post:
            _generate_server(
                prompt="test",
                count=1,
                fmt="alpaca",
                model_name="test-model",
                api_base=None,
                temperature=0.8,
                seed_examples=[],
            )

        call_headers = mock_post.call_args[1]["headers"]
        assert "Authorization" not in call_headers

    def test_generate_server_custom_api_base(self):
        """_generate_server should use custom api_base."""
        from soup_cli.commands.generate import _generate_server

        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = {
            "choices": [
                {"message": {"content": '[]'}}
            ]
        }

        with mock_patch("httpx.post", return_value=mock_response) as mock_post:
            _generate_server(
                prompt="test",
                count=1,
                fmt="alpaca",
                model_name="test-model",
                api_base="http://localhost:11434/v1",
                temperature=0.8,
                seed_examples=[],
            )

        call_url = mock_post.call_args[0][0]
        assert "localhost:11434" in call_url

    def test_generate_server_appends_v1_if_missing(self):
        """If api_base doesn't end with /v1, it should be appended."""
        from soup_cli.commands.generate import _generate_server

        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = {
            "choices": [
                {"message": {"content": '[]'}}
            ]
        }

        with mock_patch("httpx.post", return_value=mock_response) as mock_post:
            _generate_server(
                prompt="test",
                count=1,
                fmt="alpaca",
                model_name="test-model",
                api_base="http://localhost:8000",
                temperature=0.8,
                seed_examples=[],
            )

        call_url = mock_post.call_args[0][0]
        assert "/v1/chat/completions" in call_url

    def test_generate_server_error_response(self):
        """_generate_server should raise ValueError on non-200 response."""
        from soup_cli.commands.generate import _generate_server

        mock_response = MagicMock()
        mock_response.status_code = 500
        mock_response.text = "Internal Server Error"

        with mock_patch("httpx.post", return_value=mock_response):
            with pytest.raises(ValueError, match="Server returned 500"):
                _generate_server(
                    prompt="test",
                    count=1,
                    fmt="alpaca",
                    model_name="test-model",
                    api_base=None,
                    temperature=0.8,
                    seed_examples=[],
                )

    def test_generate_server_parses_json_array(self):
        """_generate_server should parse JSON array from response content."""
        from soup_cli.commands.generate import _generate_server

        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = {
            "choices": [
                {
                    "message": {
                        "content": (
                            '[{"instruction": "What is 2+2?", '
                            '"input": "", "output": "4"}]'
                        )
                    }
                }
            ]
        }

        with mock_patch("httpx.post", return_value=mock_response):
            result = _generate_server(
                prompt="test",
                count=1,
                fmt="alpaca",
                model_name="test-model",
                api_base=None,
                temperature=0.8,
                seed_examples=[],
            )

        assert len(result) == 1
        assert result[0]["instruction"] == "What is 2+2?"

    def test_generate_server_timeout(self):
        """_generate_server should use 300s timeout for local servers."""
        from soup_cli.commands.generate import _generate_server

        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = {
            "choices": [{"message": {"content": "[]"}}]
        }

        with mock_patch("httpx.post", return_value=mock_response) as mock_post:
            _generate_server(
                prompt="test",
                count=1,
                fmt="alpaca",
                model_name="test-model",
                api_base=None,
                temperature=0.8,
                seed_examples=[],
            )

        assert mock_post.call_args[1]["timeout"] == 300.0


# ─── _generate_batch Routing Tests ───────────────────────────────────────


class TestGenerateBatchRouting:
    """Test that _generate_batch routes to the correct provider."""

    def test_batch_routes_to_server(self):
        """provider='server' should route to _generate_server."""
        from soup_cli.commands.generate import _generate_batch

        with mock_patch(
            "soup_cli.commands.generate._generate_server",
            return_value=[{"instruction": "x", "output": "y"}],
        ) as mock_server:
            result = _generate_batch(
                prompt="test",
                count=1,
                fmt="alpaca",
                provider="server",
                model_name="test-model",
                api_key=None,
                api_base=None,
                temperature=0.8,
                seed_examples=[],
            )

        mock_server.assert_called_once()
        assert len(result) == 1

    def test_batch_routes_to_openai(self):
        """provider='openai' should route to _generate_openai."""
        from soup_cli.commands.generate import _generate_batch

        with mock_patch(
            "soup_cli.commands.generate._generate_openai",
            return_value=[],
        ) as mock_openai:
            _generate_batch(
                prompt="test",
                count=1,
                fmt="alpaca",
                provider="openai",
                model_name="gpt-4o-mini",
                api_key="sk-test",
                api_base=None,
                temperature=0.8,
                seed_examples=[],
            )

        mock_openai.assert_called_once()

    def test_batch_routes_to_local(self):
        """provider='local' should route to _generate_local."""
        from soup_cli.commands.generate import _generate_batch

        with mock_patch(
            "soup_cli.commands.generate._generate_local",
            return_value=[],
        ) as mock_local:
            _generate_batch(
                prompt="test",
                count=1,
                fmt="alpaca",
                provider="local",
                model_name="some-model",
                api_key=None,
                api_base=None,
                temperature=0.8,
                seed_examples=[],
            )

        mock_local.assert_called_once()


# ─── _parse_json_array Tests ─────────────────────────────────────────────


class TestParseJsonArray:
    """Test the JSON array parsing function used by all providers."""

    def test_parse_valid_json_array(self):
        """Should parse a clean JSON array."""
        from soup_cli.commands.generate import _parse_json_array

        result = _parse_json_array('[{"a": 1}, {"b": 2}]')
        assert len(result) == 2
        assert result[0]["a"] == 1

    def test_parse_markdown_code_fence(self):
        """Should strip markdown code fences."""
        from soup_cli.commands.generate import _parse_json_array

        content = '```json\n[{"instruction": "test", "output": "ok"}]\n```'
        result = _parse_json_array(content)
        assert len(result) == 1
        assert result[0]["instruction"] == "test"

    def test_parse_json_with_surrounding_text(self):
        """Should extract JSON array from surrounding text."""
        from soup_cli.commands.generate import _parse_json_array

        content = 'Here are the examples:\n[{"a": 1}]\nDone!'
        result = _parse_json_array(content)
        assert len(result) == 1

    def test_parse_ndjson_fallback(self):
        """Should fall back to line-by-line JSON parsing."""
        from soup_cli.commands.generate import _parse_json_array

        content = '{"a": 1}\n{"b": 2}\n{"c": 3}'
        result = _parse_json_array(content)
        assert len(result) == 3

    def test_parse_empty_array(self):
        """Should return empty list for empty array."""
        from soup_cli.commands.generate import _parse_json_array

        assert _parse_json_array("[]") == []

    def test_parse_invalid_json_returns_empty(self):
        """Should return empty list for completely invalid JSON."""
        from soup_cli.commands.generate import _parse_json_array

        assert _parse_json_array("not json at all") == []

    def test_parse_filters_non_dict_items(self):
        """Should filter out non-dict items from the array."""
        from soup_cli.commands.generate import _parse_json_array

        result = _parse_json_array('[{"a": 1}, "string", 42, {"b": 2}]')
        assert len(result) == 2

    def test_parse_code_fence_without_language(self):
        """Should strip code fences without language specifier."""
        from soup_cli.commands.generate import _parse_json_array

        content = '```\n[{"x": 1}]\n```'
        result = _parse_json_array(content)
        assert len(result) == 1


# ─── _validate_example Tests ─────────────────────────────────────────────


class TestValidateExample:
    """Test format validation for generated examples."""

    def test_validate_alpaca_valid(self):
        from soup_cli.commands.generate import _validate_example

        assert _validate_example({"instruction": "Q", "output": "A"}, "alpaca")

    def test_validate_alpaca_missing_output(self):
        from soup_cli.commands.generate import _validate_example

        assert not _validate_example({"instruction": "Q"}, "alpaca")

    def test_validate_sharegpt_valid(self):
        from soup_cli.commands.generate import _validate_example

        row = {
            "conversations": [
                {"from": "human", "value": "Hi"},
                {"from": "gpt", "value": "Hello"},
            ]
        }
        assert _validate_example(row, "sharegpt")

    def test_validate_sharegpt_too_few_turns(self):
        from soup_cli.commands.generate import _validate_example

        row = {"conversations": [{"from": "human", "value": "Hi"}]}
        assert not _validate_example(row, "sharegpt")

    def test_validate_chatml_valid(self):
        from soup_cli.commands.generate import _validate_example

        row = {
            "messages": [
                {"role": "user", "content": "Hi"},
                {"role": "assistant", "content": "Hello"},
            ]
        }
        assert _validate_example(row, "chatml")

    def test_validate_chatml_empty_messages(self):
        from soup_cli.commands.generate import _validate_example

        assert not _validate_example({"messages": []}, "chatml")

    def test_validate_unknown_format(self):
        from soup_cli.commands.generate import _validate_example

        assert not _validate_example({"a": 1}, "unknown")


# ─── SSRF Validation Tests ───────────────────────────────────────────────


class TestServerSSRFValidation:
    """Test SSRF protection in _generate_server."""

    def test_server_blocks_non_http_scheme(self):
        """file:// scheme should be rejected."""
        from soup_cli.commands.generate import _generate_server

        with pytest.raises(ValueError, match="HTTP or HTTPS"):
            _generate_server(
                prompt="test", count=1, fmt="alpaca",
                model_name="m", api_base="file:///etc/passwd",
                temperature=0.8, seed_examples=[],
            )

    def test_server_blocks_remote_http(self):
        """Remote HTTP (non-localhost) should be rejected."""
        from soup_cli.commands.generate import _generate_server

        with pytest.raises(ValueError, match="HTTPS for remote"):
            _generate_server(
                prompt="test", count=1, fmt="alpaca",
                model_name="m", api_base="http://169.254.169.254/latest",
                temperature=0.8, seed_examples=[],
            )

    def test_server_allows_localhost_http(self):
        """HTTP to localhost should be allowed."""
        from soup_cli.commands.generate import _generate_server

        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = {
            "choices": [{"message": {"content": "[]"}}]
        }

        with mock_patch("httpx.post", return_value=mock_response):
            result = _generate_server(
                prompt="test", count=1, fmt="alpaca",
                model_name="m", api_base="http://127.0.0.1:8000",
                temperature=0.8, seed_examples=[],
            )
        assert result == []