tests/test_llm_handler.py

"""Tests for LLM handler module."""

import os
from unittest.mock import MagicMock, patch

import pytest

from src.llm_handler import (
    DEFAULT_GROQ_MODEL,
    DEFAULT_OLLAMA_MODEL,
    GROQ_MODELS,
    OLLAMA_MODELS,
    LLMHandler,
    detect_default_provider,
    get_available_groq_models,
    get_available_ollama_models,
    get_default_groq_model,
    get_default_ollama_model,
    get_llm_handler,
)

# --- LLMHandler Tests ---


def test_init_with_default_provider_no_groq_key():
    """Test handler initialization without GROQ_API_KEY defaults to ollama."""
    with patch.dict(os.environ, {}, clear=True), patch("src.llm_handler.get_config") as mock_config:
        mock_config.return_value.get.return_value = "ollama"
        mock_config.return_value.get_env.return_value = None

        handler = LLMHandler()

        assert handler.provider == "ollama"


def test_init_with_groq_api_key_auto_detects_groq():
    """Test handler auto-detects groq provider when API key is present."""
    with patch("src.llm_handler.get_config") as mock_config:
        mock_config.return_value.get.return_value = "ollama"
        mock_config.return_value.get_env.return_value = "test-groq-api-key"

        handler = LLMHandler()

        assert handler.provider == "groq"


def test_init_with_provider_override():
    """Test handler respects provider override."""
    with patch("src.llm_handler.get_config") as mock_config:
        mock_config.return_value.get.return_value = "ollama"
        mock_config.return_value.get_env.return_value = "test-groq-api-key"

        handler = LLMHandler(provider_override="ollama")

        assert handler.provider == "ollama"


def test_init_with_api_key_override():
    """Test handler uses API key override."""
    with patch("src.llm_handler.get_config") as mock_config:
        mock_config.return_value.get.return_value = "ollama"
        mock_config.return_value.get_env.return_value = None

        handler = LLMHandler(api_key_override="custom-api-key")

        # Should detect groq because we provided an API key override
        assert handler.provider == "groq"


def test_get_groq_api_key_from_override():
    """Test getting Groq API key from override."""
    with patch("src.llm_handler.get_config") as mock_config:
        mock_config.return_value.get.return_value = "groq"
        mock_config.return_value.get_env.return_value = None

        handler = LLMHandler(api_key_override="override-key")

        assert handler._get_groq_api_key() == "override-key"


def test_get_groq_api_key_from_env():
    """Test getting Groq API key from environment."""
    with patch("src.llm_handler.get_config") as mock_config:
        mock_config.return_value.get.return_value = "groq"
        mock_config.return_value.get_env.return_value = "env-groq-key"

        handler = LLMHandler(provider_override="groq")

        assert handler._get_groq_api_key() == "env-groq-key"


def test_get_groq_llm_without_api_key_raises():
    """Test get_groq_llm raises ValueError without API key."""
    with patch("src.llm_handler.get_config") as mock_config:
        mock_config.return_value.get.return_value = "groq"
        mock_config.return_value.get_env.return_value = None

        handler = LLMHandler(provider_override="groq")

        with pytest.raises(ValueError, match="GROQ_API_KEY not found"):
            handler.get_groq_llm()


def test_get_groq_llm_with_api_key():
    """Test get_groq_llm initializes with API key."""
    with patch("src.llm_handler.get_config") as mock_config:
        mock_config.return_value.get.side_effect = lambda key, default=None: {
            "llm.provider": "groq",
            "llm.groq_model": "llama-3.3-70b-versatile",
            "llm.temperature": 0.7,
            "llm.max_tokens": 800,
        }.get(key, default)
        mock_config.return_value.get_env.return_value = None

        handler = LLMHandler(provider_override="groq", api_key_override="test-api-key")

        # ChatGroq is imported inside the function, so patch in langchain_groq
        with patch("langchain_groq.ChatGroq") as mock_groq:
            mock_llm = MagicMock()
            mock_llm.invoke.return_value = MagicMock(content="Test response")
            mock_groq.return_value = mock_llm

            llm = handler.get_groq_llm()

            assert llm is not None
            mock_groq.assert_called_once()


def test_get_groq_llm_connection_error():
    """Test get_groq_llm handles connection errors."""
    with patch("src.llm_handler.get_config") as mock_config:
        mock_config.return_value.get.side_effect = lambda key, default=None: {
            "llm.provider": "groq",
            "llm.groq_model": "llama-3.3-70b-versatile",
            "llm.temperature": 0.7,
            "llm.max_tokens": 800,
        }.get(key, default)
        mock_config.return_value.get_env.return_value = None

        handler = LLMHandler(provider_override="groq", api_key_override="test-api-key")

        # ChatGroq is imported inside the function
        with patch("langchain_groq.ChatGroq") as mock_groq:
            mock_groq.side_effect = Exception("Connection error")

            with pytest.raises(Exception, match="Connection error"):
                handler.get_groq_llm()


def test_get_ollama_llm():
    """Test get_ollama_llm initialization."""
    with patch("src.llm_handler.get_config") as mock_config:
        mock_config.return_value.get.side_effect = lambda key, default=None: {
            "llm.provider": "ollama",
            "llm.model": "llama3.2:3b",
            "llm.temperature": 0.7,
            "llm.max_tokens": 512,
            "llm.top_p": 0.9,
        }.get(key, default)
        mock_config.return_value.get_env.side_effect = lambda key, default=None: {
            "OLLAMA_BASE_URL": "http://localhost:11434",
        }.get(key, default)

        handler = LLMHandler(provider_override="ollama")

        with patch("src.llm_handler.OllamaLLM") as mock_ollama:
            mock_llm = MagicMock()
            mock_llm.invoke.return_value = "Test response from Ollama"
            mock_ollama.return_value = mock_llm

            llm = handler.get_ollama_llm()

            assert llm is not None
            mock_ollama.assert_called_once()


def test_get_ollama_llm_model_not_found():
    """Test get_ollama_llm raises helpful error for missing model."""
    with patch("src.llm_handler.get_config") as mock_config:
        mock_config.return_value.get.side_effect = lambda key, default=None: {
            "llm.provider": "ollama",
            "llm.model": "nonexistent-model",
            "llm.temperature": 0.7,
            "llm.max_tokens": 512,
            "llm.top_p": 0.9,
        }.get(key, default)
        mock_config.return_value.get_env.side_effect = lambda key, default=None: {
            "OLLAMA_BASE_URL": "http://localhost:11434",
        }.get(key, default)

        handler = LLMHandler(provider_override="ollama")

        with patch("src.llm_handler.OllamaLLM") as mock_ollama:
            mock_ollama.return_value.invoke.side_effect = Exception(
                "model not found: nonexistent-model"
            )

            with pytest.raises(ValueError, match="not found in Ollama"):
                handler.get_ollama_llm()


def test_get_ollama_llm_other_error():
    """Test get_ollama_llm re-raises non-model-not-found errors."""
    with patch("src.llm_handler.get_config") as mock_config:
        mock_config.return_value.get.side_effect = lambda key, default=None: {
            "llm.provider": "ollama",
            "llm.model": "llama3.2:3b",
            "llm.temperature": 0.7,
            "llm.max_tokens": 512,
            "llm.top_p": 0.9,
        }.get(key, default)
        mock_config.return_value.get_env.side_effect = lambda key, default=None: {
            "OLLAMA_BASE_URL": "http://localhost:11434",
        }.get(key, default)

        handler = LLMHandler(provider_override="ollama")

        with patch("src.llm_handler.OllamaLLM") as mock_ollama:
            mock_ollama.return_value.invoke.side_effect = Exception("Connection refused")

            with pytest.raises(Exception, match="Connection refused"):
                handler.get_ollama_llm()


def test_get_llm_returns_cached_instance():
    """Test get_llm returns cached LLM instance."""
    with patch("src.llm_handler.get_config") as mock_config:
        mock_config.return_value.get.return_value = "ollama"
        mock_config.return_value.get_env.return_value = None

        handler = LLMHandler(provider_override="ollama")

        # Set a pre-cached LLM
        mock_llm = MagicMock()
        handler.llm = mock_llm

        result = handler.get_llm()

        assert result is mock_llm


def test_get_llm_unsupported_provider_raises():
    """Test get_llm raises for unsupported provider."""
    with patch("src.llm_handler.get_config") as mock_config:
        mock_config.return_value.get.return_value = "unsupported"
        mock_config.return_value.get_env.return_value = None

        handler = LLMHandler(provider_override="unsupported")

        with pytest.raises(ValueError, match="Unsupported LLM provider"):
            handler.get_llm()


def test_get_system_prompt():
    """Test get_system_prompt returns formatted prompt."""
    with patch("src.llm_handler.get_config") as mock_config:
        mock_config.return_value.get.side_effect = lambda key, default=None: {
            "llm.provider": "ollama",
            "llm.system_prompt": "Hello {name}, you are a {title}",
        }.get(key, default)
        mock_config.return_value.get_env.return_value = None
        mock_config.return_value.format_template.return_value = "Hello Tin, you are a Engineer"

        handler = LLMHandler(provider_override="ollama")
        prompt = handler.get_system_prompt()

        assert prompt == "Hello Tin, you are a Engineer"


def test_get_provider():
    """Test get_provider returns current provider."""
    with patch("src.llm_handler.get_config") as mock_config:
        mock_config.return_value.get.return_value = "ollama"
        mock_config.return_value.get_env.return_value = None

        handler = LLMHandler(provider_override="groq")

        assert handler.get_provider() == "groq"


def test_get_model_with_override():
    """Test get_model returns override when provided."""
    with patch("src.llm_handler.get_config") as mock_config:
        mock_config.return_value.get.return_value = "ollama"
        mock_config.return_value.get_env.return_value = None

        handler = LLMHandler(provider_override="ollama", model_override="custom-model")

        assert handler.get_model() == "custom-model"


def test_get_model_groq_provider():
    """Test get_model returns groq model for groq provider."""
    with patch("src.llm_handler.get_config") as mock_config:
        mock_config.return_value.get.side_effect = lambda key, default=None: {
            "llm.provider": "groq",
            "llm.groq_model": "llama-3.3-70b-versatile",
        }.get(key, default)
        mock_config.return_value.get_env.return_value = "test-key"

        handler = LLMHandler(provider_override="groq")

        assert handler.get_model() == "llama-3.3-70b-versatile"


def test_get_model_ollama_provider():
    """Test get_model returns ollama model for ollama provider."""
    with patch("src.llm_handler.get_config") as mock_config:
        mock_config.return_value.get.side_effect = lambda key, default=None: {
            "llm.provider": "ollama",
            "llm.model": "llama3.2:3b",
        }.get(key, default)
        mock_config.return_value.get_env.return_value = None

        handler = LLMHandler(provider_override="ollama")

        assert handler.get_model() == "llama3.2:3b"


# --- Helper Functions Tests ---


def test_get_llm_handler():
    """Test get_llm_handler factory function."""
    with patch("src.llm_handler.get_config") as mock_config:
        mock_config.return_value.get.return_value = "ollama"
        mock_config.return_value.get_env.return_value = None

        handler = get_llm_handler(provider_override="ollama")

        assert isinstance(handler, LLMHandler)


def test_get_available_groq_models():
    """Test get_available_groq_models returns list copy."""
    models = get_available_groq_models()

    assert isinstance(models, list)
    assert len(models) > 0
    assert models == GROQ_MODELS
    # Verify it's a copy
    models.append("test-model")
    assert "test-model" not in GROQ_MODELS


def test_get_default_groq_model():
    """Test get_default_groq_model returns default."""
    model = get_default_groq_model()

    assert model == DEFAULT_GROQ_MODEL


def test_get_available_ollama_models_success():
    """Test get_available_ollama_models fetches from API."""
    with patch("src.llm_handler.get_config") as mock_config:
        mock_config.return_value.get_env.return_value = "http://localhost:11434"

        # ollama is imported inside the function
        with patch.dict("sys.modules", {"ollama": MagicMock()}):
            import sys

            mock_ollama = sys.modules["ollama"]
            mock_model1 = MagicMock()
            mock_model1.model = "llama3.2:3b"
            mock_model2 = MagicMock()
            mock_model2.model = "phi3:mini"

            mock_response = MagicMock()
            mock_response.models = [mock_model1, mock_model2]
            mock_ollama.Client.return_value.list.return_value = mock_response

            models = get_available_ollama_models()

            assert "llama3.2:3b" in models
            assert "phi3:mini" in models


def test_get_available_ollama_models_empty():
    """Test get_available_ollama_models returns empty list when no models."""
    with patch("src.llm_handler.get_config") as mock_config:
        mock_config.return_value.get_env.return_value = "http://localhost:11434"

        # ollama is imported inside the function
        with patch.dict("sys.modules", {"ollama": MagicMock()}):
            import sys

            mock_ollama = sys.modules["ollama"]
            mock_response = MagicMock()
            mock_response.models = []
            mock_ollama.Client.return_value.list.return_value = mock_response

            models = get_available_ollama_models()

            assert models == []


def test_get_available_ollama_models_fallback_on_error():
    """Test get_available_ollama_models falls back to static list on error."""
    with patch("src.llm_handler.get_config") as mock_config:
        mock_config.return_value.get_env.return_value = "http://localhost:11434"

        # ollama is imported inside the function
        with patch.dict("sys.modules", {"ollama": MagicMock()}):
            import sys

            mock_ollama = sys.modules["ollama"]
            mock_ollama.Client.side_effect = Exception("Connection failed")

            models = get_available_ollama_models()

            assert models == OLLAMA_MODELS


def test_get_default_ollama_model():
    """Test get_default_ollama_model returns default."""
    model = get_default_ollama_model()

    assert model == DEFAULT_OLLAMA_MODEL


def test_detect_default_provider_with_groq_key():
    """Test detect_default_provider returns groq when key present."""
    # Path is imported from pathlib inside the function
    with patch("pathlib.Path.exists", return_value=False):
        with patch.dict(os.environ, {"GROQ_API_KEY": "test-key"}):
            provider = detect_default_provider()

            assert provider == "groq"


def test_detect_default_provider_without_groq_key():
    """Test detect_default_provider returns ollama when no key."""
    # Path is imported from pathlib inside the function
    with patch("pathlib.Path.exists", return_value=False):
        with patch.dict(os.environ, {}, clear=True):
            # Remove GROQ_API_KEY if it exists
            os.environ.pop("GROQ_API_KEY", None)
            provider = detect_default_provider()

            assert provider == "ollama"