zzrouter/backend/tests/test_litellm_wrapper.py

"""Tests for LiteLLM wrapper."""
import pytest
from unittest.mock import patch, AsyncMock, MagicMock
from app.providers.litellm_wrapper import LiteLLMWrapper
from app.models import Provider


@pytest.fixture
def openai_provider():
    """Return an OpenAI provider."""
    return Provider(
        id=1,
        name="openai",
        base_url="https://api.openai.com/v1",
        api_type="openai",
        is_active=True
    )


@pytest.fixture
def anthropic_provider():
    """Return an Anthropic provider."""
    return Provider(
        id=2,
        name="anthropic",
        base_url="https://api.anthropic.com/v1",
        api_type="anthropic",
        is_active=True
    )


class TestBuildModelString:
    """Test _build_model_string method."""

    def test_openai_model(self, openai_provider):
        """Test building model string for OpenAI."""
        result = LiteLLMWrapper._build_model_string(openai_provider, "gpt-4o")
        assert result == "openai/gpt-4o"

    def test_anthropic_model(self, anthropic_provider):
        """Test building model string for Anthropic."""
        result = LiteLLMWrapper._build_model_string(anthropic_provider, "claude-3-opus")
        assert result == "anthropic/claude-3-opus"

    def test_unknown_provider_type(self):
        """Test with unknown provider type."""
        provider = Provider(
            id=3,
            name="unknown",
            base_url="https://api.example.com",
            api_type="unknown",
            is_active=True
        )
        result = LiteLLMWrapper._build_model_string(provider, "model-name")
        assert result == "openai/model-name"  # Falls back to openai


class TestChatCompletion:
    """Test chat_completion method."""

    @pytest.mark.asyncio
    async def test_chat_completion_non_streaming(self, openai_provider):
        """Test non-streaming chat completion."""
        mock_response = MagicMock()
        mock_response.usage = MagicMock()
        mock_response.usage.prompt_tokens = 15
        mock_response.usage.completion_tokens = 25

        with patch("app.providers.litellm_wrapper.acompletion", new_callable=AsyncMock) as mock_acompletion:
            mock_acompletion.return_value = mock_response

            response, prompt_tokens, completion_tokens = await LiteLLMWrapper.chat_completion(
                provider=openai_provider,
                api_key="sk-test",
                model="gpt-4o",
                messages=[{"role": "user", "content": "hi"}],
                stream=False
            )

            assert response == mock_response
            assert prompt_tokens == 15
            assert completion_tokens == 25

    @pytest.mark.asyncio
    async def test_chat_completion_streaming(self, openai_provider):
        """Test streaming chat completion."""
        mock_stream = MagicMock()

        with patch("app.providers.litellm_wrapper.acompletion", new_callable=AsyncMock) as mock_acompletion:
            mock_acompletion.return_value = mock_stream

            response, prompt_tokens, completion_tokens = await LiteLLMWrapper.chat_completion(
                provider=openai_provider,
                api_key="sk-test",
                model="gpt-4o",
                messages=[{"role": "user", "content": "hi"}],
                stream=True
            )

            # For streaming, token counts are 0 initially
            assert prompt_tokens == 0
            assert completion_tokens == 0

    @pytest.mark.asyncio
    async def test_chat_completion_with_kwargs(self, openai_provider):
        """Test chat completion with additional kwargs."""
        mock_response = MagicMock()
        mock_response.usage = MagicMock()
        mock_response.usage.prompt_tokens = 10
        mock_response.usage.completion_tokens = 20

        with patch("app.providers.litellm_wrapper.acompletion", new_callable=AsyncMock) as mock_acompletion:
            mock_acompletion.return_value = mock_response

            await LiteLLMWrapper.chat_completion(
                provider=openai_provider,
                api_key="sk-test",
                model="gpt-4o",
                messages=[{"role": "user", "content": "hi"}],
                stream=False,
                temperature=0.7,
                max_tokens=100
            )

            # Verify kwargs were passed
            call_kwargs = mock_acompletion.call_args[1]
            assert call_kwargs["temperature"] == 0.7
            assert call_kwargs["max_tokens"] == 100

    @pytest.mark.asyncio
    async def test_chat_completion_anthropic(self, anthropic_provider):
        """Test chat completion with Anthropic provider."""
        mock_response = MagicMock()
        mock_response.usage = MagicMock()
        mock_response.usage.prompt_tokens = 10
        mock_response.usage.completion_tokens = 30

        with patch("app.providers.litellm_wrapper.acompletion", new_callable=AsyncMock) as mock_acompletion:
            mock_acompletion.return_value = mock_response

            response, prompt_tokens, completion_tokens = await LiteLLMWrapper.chat_completion(
                provider=anthropic_provider,
                api_key="sk-ant-test",
                model="claude-3-opus",
                messages=[{"role": "user", "content": "hi"}],
                stream=False
            )

            # Verify model string was built correctly
            call_kwargs = mock_acompletion.call_args[1]
            assert call_kwargs["model"] == "anthropic/claude-3-opus"


class TestStreamResponse:
    """Test stream_response method."""

    @pytest.mark.asyncio
    async def test_stream_response_yields_chunks(self, openai_provider):
        """Test stream_response yields SSE-formatted chunks."""
        mock_chunk1 = MagicMock()
        mock_chunk1.model_dump.return_value = {"choices": [{"delta": {"content": "Hello"}}]}
        mock_chunk1.usage = None

        mock_chunk2 = MagicMock()
        mock_chunk2.model_dump.return_value = {"choices": [{"delta": {"content": "!"}}]}
        mock_chunk2.usage = MagicMock(prompt_tokens=5, completion_tokens=2)

        async def mock_stream():
            yield mock_chunk1
            yield mock_chunk2

        with patch("app.providers.litellm_wrapper.async_logger.log", new_callable=AsyncMock):
            chunks = []
            async for chunk in LiteLLMWrapper.stream_response(
                mock_stream(),
                client_key_id=1,
                provider_id=1,
                model="gpt-4o",
                start_time=0
            ):
                chunks.append(chunk)

            assert len(chunks) >= 2  # At least 2 chunks + [DONE]
            assert any("Hello" in chunk for chunk in chunks)
            assert "data: [DONE]" in chunks[-1]

    @pytest.mark.asyncio
    async def test_stream_response_handles_error(self, openai_provider):
        """Test stream_response handles errors."""
        async def mock_stream_with_error():
            yield MagicMock(model_dump=lambda: {"choices": []})
            raise Exception("Stream error")

        with patch("app.providers.litellm_wrapper.async_logger.log", new_callable=AsyncMock):
            chunks = []
            async for chunk in LiteLLMWrapper.stream_response(
                mock_stream_with_error(),
                client_key_id=1,
                provider_id=1,
                model="gpt-4o",
                start_time=0
            ):
                chunks.append(chunk)

            # Should include error message
            assert any("error" in chunk.lower() for chunk in chunks)