office_translator/tests/test_providers/test_ollama_provider.py

"""
Tests for the OllamaTranslationProvider.
"""

import pytest
from unittest.mock import patch, MagicMock
from requests.exceptions import Timeout, ConnectionError as RequestsConnectionError

from services.providers.ollama_provider import (
    OllamaTranslationProvider,
    OllamaProviderError,
    get_ollama_provider,
    register_ollama_provider,
    _build_system_prompt,
    _get_language_name,
    OLLAMA_UNAVAILABLE,
    OLLAMA_MODEL_NOT_FOUND,
    OLLAMA_TIMEOUT,
    OLLAMA_GENERATION_ERROR,
    OLLAMA_CONTEXT_TOO_LONG,
)
from services.providers.schemas import TranslationRequest, TranslationResponse


class TestOllamaProviderError:
    """Tests for OllamaProviderError exception."""

    def test_error_creation(self):
        """Test error creation with all fields."""
        error = OllamaProviderError(
            code=OLLAMA_UNAVAILABLE,
            message="Ollama unavailable",
            details={"provider": "ollama"},
        )

        assert error.code == OLLAMA_UNAVAILABLE
        assert error.message == "Ollama unavailable"
        assert error.details == {"provider": "ollama"}

    def test_error_to_dict(self):
        """Test error serialization."""
        error = OllamaProviderError(
            code=OLLAMA_MODEL_NOT_FOUND,
            message="Model not found",
            details={"model": "llama3"},
        )

        result = error.to_dict()

        assert result["error"] == OLLAMA_MODEL_NOT_FOUND
        assert result["message"] == "Model not found"
        assert result["details"]["model"] == "llama3"

    def test_error_to_dict_no_details(self):
        """Test error serialization without details."""
        error = OllamaProviderError(
            code=OLLAMA_GENERATION_ERROR,
            message="Generation error",
        )

        result = error.to_dict()

        assert result["error"] == OLLAMA_GENERATION_ERROR
        assert result["message"] == "Generation error"
        assert "details" not in result


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

    def test_get_language_name_common(self):
        """Test language name lookup for common languages."""
        assert _get_language_name("en") == "English"
        assert _get_language_name("fr") == "French"
        assert _get_language_name("es") == "Spanish"
        assert _get_language_name("de") == "German"
        assert _get_language_name("zh") == "Chinese"
        assert _get_language_name("ja") == "Japanese"

    def test_get_language_name_with_variant(self):
        """Test language name lookup with variant codes."""
        assert _get_language_name("en-US") == "English"
        assert _get_language_name("pt-BR") == "Portuguese"

    def test_get_language_name_unknown(self):
        """Test language name lookup for unknown codes."""
        assert _get_language_name("xx") == "xx"

    def test_build_system_prompt_default(self):
        """Test default system prompt generation."""
        prompt = _build_system_prompt("English", "French")

        assert "English" in prompt
        assert "French" in prompt
        assert "translator" in prompt.lower()

    def test_build_system_prompt_custom(self):
        """Test custom system prompt."""
        custom = "Translate this text formally for business context."
        prompt = _build_system_prompt("English", "French", custom)

        assert prompt == custom


class TestOllamaTranslationProvider:
    """Tests for OllamaTranslationProvider."""

    @pytest.fixture
    def provider(self):
        """Create an Ollama provider instance."""
        return OllamaTranslationProvider(
            base_url="http://localhost:11434",
            model="llama3",
            timeout=120,
            max_retries=0,
        )

    @pytest.fixture
    def provider_with_retries(self):
        """Create an Ollama provider with retries."""
        return OllamaTranslationProvider(
            base_url="http://localhost:11434",
            model="llama3",
            timeout=120,
            max_retries=2,
            retry_delay=0.01,
        )

    def test_init(self, provider):
        """Test provider initialization."""
        assert provider._base_url == "http://localhost:11434"
        assert provider._model == "llama3"
        assert provider.timeout == 120
        assert provider._provider_name == "ollama"

    def test_get_name(self, provider):
        """Test provider name."""
        assert provider.get_name() == "ollama"

    def test_translate_text_empty(self, provider):
        """Test translating empty text."""
        request = TranslationRequest(text="", target_language="fr")
        response = provider.translate_text(request)

        assert response.translated_text == ""
        assert response.provider_name == "ollama"
        assert response.from_cache is False

    def test_translate_text_whitespace(self, provider):
        """Test translating whitespace-only text."""
        request = TranslationRequest(text="   ", target_language="fr")
        response = provider.translate_text(request)

        assert response.translated_text == "   "

    @patch.object(OllamaTranslationProvider, "_fetch_available_models")
    @patch.object(OllamaTranslationProvider, "_make_api_request")
    def test_translate_text_success(self, mock_request, mock_models, provider):
        """Test successful translation."""
        mock_models.return_value = ["llama3", "mistral"]
        mock_request.return_value = "Bonjour"

        request = TranslationRequest(text="Hello", target_language="fr")
        response = provider.translate_text(request)

        assert response.translated_text == "Bonjour"
        assert response.provider_name == "ollama"
        assert response.from_cache is False

    @patch.object(OllamaTranslationProvider, "_fetch_available_models")
    @patch.object(OllamaTranslationProvider, "_make_api_request")
    def test_translate_text_with_custom_prompt(
        self, mock_request, mock_models, provider
    ):
        """Test translation with custom system prompt."""
        mock_models.return_value = ["llama3"]
        mock_request.return_value = "Bonjour (formal)"

        request = TranslationRequest(
            text="Hello",
            target_language="fr",
            metadata={"custom_prompt": "Translate formally for business"},
        )
        response = provider.translate_text(request)

        assert response.translated_text == "Bonjour (formal)"
        mock_request.assert_called_once()
        call_args = mock_request.call_args
        assert "Translate formally for business" in call_args[0][1]

    def test_translate_batch_empty(self, provider):
        """Test batch translation with empty list."""
        responses = provider.translate_batch([])
        assert responses == []

    @patch.object(OllamaTranslationProvider, "translate_text")
    def test_translate_batch(self, mock_translate, provider):
        """Test batch translation."""
        mock_translate.side_effect = [
            TranslationResponse(translated_text="Bonjour", provider_name="ollama"),
            TranslationResponse(translated_text="Monde", provider_name="ollama"),
        ]

        requests = [
            TranslationRequest(text="Hello", target_language="fr"),
            TranslationRequest(text="World", target_language="fr"),
        ]
        responses = provider.translate_batch(requests)

        assert len(responses) == 2
        assert responses[0].translated_text == "Bonjour"
        assert responses[1].translated_text == "Monde"


class TestOllamaErrorCodes:
    """Tests for Ollama error code handling."""

    @pytest.fixture
    def provider(self):
        """Create an Ollama provider with no retries for error testing."""
        return OllamaTranslationProvider(
            base_url="http://localhost:11434",
            model="llama3",
            timeout=120,
            max_retries=0,
        )

    def test_context_too_long_error(self, provider):
        """Test context too long error."""
        long_text = "x" * 130000
        request = TranslationRequest(text=long_text, target_language="fr")
        response = provider.translate_text(request)

        assert response.error_code == OLLAMA_CONTEXT_TOO_LONG
        assert response.error is not None

    @patch.object(OllamaTranslationProvider, "_fetch_available_models")
    def test_model_not_found_error(self, mock_models, provider):
        """Test model not found error."""
        mock_models.return_value = ["mistral", "qwen2"]

        request = TranslationRequest(text="Hello", target_language="fr")
        response = provider.translate_text(request)

        assert response.error_code == OLLAMA_MODEL_NOT_FOUND
        assert "llama3" in response.error

    @patch.object(OllamaTranslationProvider, "_check_model_available")
    @patch("services.providers.ollama_provider.requests")
    def test_unavailable_error(self, mock_requests, mock_model_available, provider):
        """Test Ollama unavailable error."""
        mock_model_available.return_value = True
        mock_requests.post.side_effect = RequestsConnectionError("Connection refused")

        request = TranslationRequest(text="Hello", target_language="fr")
        response = provider.translate_text(request)

        assert response.error_code == OLLAMA_UNAVAILABLE
        assert "indisponible" in response.error.lower()


class TestOllamaHealthCheck:
    """Tests for Ollama health check functionality."""

    @pytest.fixture
    def provider(self):
        """Create an Ollama provider instance."""
        return OllamaTranslationProvider(
            base_url="http://localhost:11434",
            model="llama3",
        )

    @patch("services.providers.ollama_provider.requests")
    def test_health_check_available(self, mock_requests, provider):
        """Test health check when Ollama is available."""
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = {
            "models": [{"name": "llama3"}, {"name": "mistral"}]
        }
        mock_requests.get.return_value = mock_response

        status = provider.health_check()

        assert status.name == "ollama"
        assert status.available is True
        assert status.latency_ms is not None
        assert status.model == "llama3"
        assert status.model_available is True

    @patch("services.providers.ollama_provider.requests")
    def test_health_check_model_not_pulled(self, mock_requests, provider):
        """Test health check when model is not pulled."""
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = {"models": [{"name": "mistral"}]}
        mock_requests.get.return_value = mock_response

        status = provider.health_check()

        assert status.available is False
        assert "llama3" in status.error
        assert status.model == "llama3"
        assert status.model_available is False

    @patch("services.providers.ollama_provider.requests")
    def test_health_check_unavailable(self, mock_requests, provider):
        """Test health check when Ollama is unavailable."""
        mock_requests.get.side_effect = RequestsConnectionError("Connection refused")

        status = provider.health_check()

        assert status.available is False

    @patch("services.providers.ollama_provider.requests")
    def test_health_check_caching(self, mock_requests, provider):
        """Test that health check results are cached (no API call when cache valid)."""
        import time
        from services.providers.schemas import ProviderHealthStatus

        current_time = time.time()
        cached_status = ProviderHealthStatus(
            name="ollama",
            available=True,
            latency_ms=50.0,
            error=None,
            last_check="2024-01-15T10:00:00Z",
        )
        provider._health_cache["health_check"] = {
            "value": cached_status,
            "timestamp": current_time,
        }

        status = provider.health_check()

        assert status.available is True
        mock_requests.get.assert_not_called()


class TestOllamaProviderRetry:
    """Tests for Ollama provider retry logic."""

    @pytest.fixture
    def provider(self):
        """Create an Ollama provider with retry enabled."""
        return OllamaTranslationProvider(
            base_url="http://localhost:11434",
            model="llama3",
            max_retries=2,
            retry_delay=0.01,
        )

    @patch.object(OllamaTranslationProvider, "_fetch_available_models")
    @patch.object(OllamaTranslationProvider, "_make_api_request")
    def test_retry_on_timeout(self, mock_request, mock_models, provider):
        """Test that timeout errors trigger retry."""
        mock_models.return_value = ["llama3"]
        mock_request.side_effect = [
            OllamaProviderError(OLLAMA_TIMEOUT, "Timeout"),
            "Bonjour",
        ]

        request = TranslationRequest(text="Hello", target_language="fr")
        response = provider.translate_text(request)

        assert response.translated_text == "Bonjour"
        assert mock_request.call_count == 2

    @patch.object(OllamaTranslationProvider, "_fetch_available_models")
    @patch.object(OllamaTranslationProvider, "_make_api_request")
    def test_no_retry_on_model_not_found(self, mock_request, mock_models, provider):
        """Test that model not found errors do not trigger retry."""
        mock_models.return_value = ["llama3"]
        mock_request.side_effect = OllamaProviderError(
            OLLAMA_MODEL_NOT_FOUND, "Model not found"
        )

        request = TranslationRequest(text="Hello", target_language="fr")
        provider.translate_text(request)

        assert mock_request.call_count == 1

    @patch.object(OllamaTranslationProvider, "_fetch_available_models")
    @patch.object(OllamaTranslationProvider, "_make_api_request")
    def test_timeout_returns_ollama_timeout_error(self, mock_request, mock_models):
        """Test that timeout without retry returns OLLAMA_TIMEOUT in response."""
        provider = OllamaTranslationProvider(
            base_url="http://localhost:11434",
            model="llama3",
            timeout=120,
            max_retries=0,
        )
        mock_models.return_value = ["llama3"]
        mock_request.side_effect = OllamaProviderError(
            OLLAMA_TIMEOUT,
            "Délai d'attente Ollama dépassé. Réessayez avec un texte plus court.",
        )

        request = TranslationRequest(text="Hello", target_language="fr")
        response = provider.translate_text(request)

        assert response.error_code == OLLAMA_TIMEOUT
        assert response.error is not None
        assert "Délai" in response.error or "timeout" in response.error.lower()


class TestOllamaProviderSingleton:
    """Tests for Ollama provider singleton functions."""

    def test_get_ollama_provider(self):
        """Test get_ollama_provider creates instance with config."""
        import services.providers.ollama_provider as ollama_module

        ollama_module._provider_instance = None

        with patch(
            "services.providers.config.ProvidersConfig"
        ) as mock_config:
            mock_config.OLLAMA_BASE_URL = "http://localhost:11434"
            mock_config.OLLAMA_MODEL = "llama3"
            mock_config.OLLAMA_TIMEOUT = 120
            mock_config.OLLAMA_MAX_RETRIES = 2
            mock_config.OLLAMA_RETRY_DELAY = 2.0

            provider = ollama_module.get_ollama_provider()

            assert provider is not None
            assert provider._model == "llama3"

            ollama_module._provider_instance = None


class TestOllamaRegistryIntegration:
    """Tests for Ollama provider registry integration."""

    def test_register_ollama_provider(self):
        """Test provider registration."""
        from services.providers.registry import registry

        registry.unregister("ollama")

        with patch(
            "services.providers.ollama_provider.get_ollama_provider"
        ) as mock_get:
            mock_provider = MagicMock()
            mock_get.return_value = mock_provider

            result = register_ollama_provider()

            assert result == mock_provider
            assert "ollama" in registry
            registry.unregister("ollama")


class TestOllamaModelCheck:
    """Tests for Ollama model availability checking."""

    @pytest.fixture
    def provider(self):
        """Create an Ollama provider instance."""
        return OllamaTranslationProvider(
            base_url="http://localhost:11434",
            model="llama3",
        )

    @patch("services.providers.ollama_provider.requests")
    def test_fetch_available_models(self, mock_requests, provider):
        """Test fetching available models."""
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = {
            "models": [
                {"name": "llama3:latest"},
                {"name": "mistral:latest"},
            ]
        }
        mock_requests.get.return_value = mock_response

        models = provider._fetch_available_models()

        assert "llama3:latest" in models
        assert "mistral:latest" in models

    def test_check_model_available(self, provider):
        """Test model availability checking."""
        import time

        provider._available_models = ["llama3:latest", "mistral:latest"]
        provider._models_cache_time = time.time()

        assert provider._check_model_available("llama3") is True
        assert provider._check_model_available("mistral") is True
        assert provider._check_model_available("qwen2") is False