python/tests/test_hybrid_routing.py

from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from src.server.services.credential_service import credential_service
from src.server.services.llm.clients import get_llm_client
from src.server.services.llm.hybrid_router import hybrid_router


def dummy_retry_decorator(*args, **kwargs):
    def decorator(func):
        return func
    return decorator

@pytest.fixture(autouse=True)
def mock_retry():
    with patch("src.server.utils.retry_utils.retry_with_backoff", dummy_retry_decorator):
        yield

@pytest.fixture(autouse=True)
def reset_active_tier():
    credential_service.set_active_tier(1)
    yield
    credential_service.set_active_tier(1)

@pytest.fixture
def mock_provider_config():
    return {
        "provider": "google",
        "api_key": "test-key",
        "base_url": "https://generativelanguage.googleapis.com/v1beta/openai/",
        "chat_model": "gemini-3.1-flash",
        "embedding_model": "text-embedding-004",
    }

@pytest.mark.asyncio
@patch("openai.AsyncOpenAI")
async def test_simple_offline_query_routes_to_tier3(mock_openai, mock_provider_config):
    """Test that a simple offline-compatible query routes to Tier 3 (Ollama) when available."""
    # Mock Ollama capability to be available
    mock_matrix = {
        "models": {
            "gemma3:4b": {"available": True, "tokens_per_sec": 4.5}
        }
    }

    # Mock Tier 3 Client (Ollama)
    mock_ollama_client = MagicMock()
    mock_ollama_client.close = AsyncMock()
    mock_ollama_client.aclose = AsyncMock()
    mock_ollama_response = MagicMock()
    mock_ollama_response.choices = [MagicMock(message=MagicMock(content="Local Ollama Response"))]
    mock_ollama_client.chat.completions.create = AsyncMock(return_value=mock_ollama_response)

    # First client setup is for Tier 3
    mock_openai.return_value = mock_ollama_client

    with patch.object(credential_service, "get_active_provider", AsyncMock(return_value=mock_provider_config)), \
         patch.object(credential_service, "get_credential", AsyncMock(return_value="0")), \
         patch.object(credential_service, "set_active_tier") as mock_set_tier, \
         patch.object(hybrid_router, "capability_matrix", mock_matrix):

        async with get_llm_client() as client:
            res = await client.chat.completions.create(messages=[{"role": "user", "content": "Hello, how are you today?"}])
            assert res.choices[0].message.content == "Local Ollama Response"
            mock_set_tier.assert_called_with(3)

@pytest.mark.asyncio
@patch("openai.AsyncOpenAI")
async def test_complex_online_query_routes_to_tier1(mock_openai, mock_provider_config):
    """Test that a query containing search/online keywords or with long length routes to Tier 1."""
    # Mock Ollama capability to be available
    mock_matrix = {
        "models": {
            "gemma3:4b": {"available": True, "tokens_per_sec": 4.5}
        }
    }

    # Mock Tier 1 Client (Cloud)
    mock_cloud_client = MagicMock()
    mock_cloud_client.close = AsyncMock()
    mock_cloud_client.aclose = AsyncMock()
    mock_cloud_response = MagicMock()
    mock_cloud_response.choices = [MagicMock(message=MagicMock(content="Cloud Response"))]
    mock_cloud_client.chat.completions.create = AsyncMock(return_value=mock_cloud_response)

    mock_openai.return_value = mock_cloud_client

    with patch.object(credential_service, "get_active_provider", AsyncMock(return_value=mock_provider_config)), \
         patch.object(credential_service, "get_credential", AsyncMock(return_value="0")), \
         patch.object(credential_service, "set_active_tier") as mock_set_tier, \
         patch.object(hybrid_router, "capability_matrix", mock_matrix):

        async with get_llm_client() as client:
            # Query containing "search" (online keyword)
            res = await client.chat.completions.create(messages=[{"role": "user", "content": "search for the latest python updates"}])
            assert res.choices[0].message.content == "Cloud Response"
            mock_set_tier.assert_called_with(1)

@pytest.mark.asyncio
@patch("openai.AsyncOpenAI")
async def test_ollama_unavailable_routes_to_tier1(mock_openai, mock_provider_config):
    """Test that a simple query routes to Tier 1 when Ollama is marked as unavailable."""
    # Mock Ollama capability to be unavailable
    mock_matrix = {
        "models": {
            "gemma3:4b": {"available": False, "tokens_per_sec": 4.5}
        }
    }

    # Mock Tier 1 Client (Cloud)
    mock_cloud_client = MagicMock()
    mock_cloud_client.close = AsyncMock()
    mock_cloud_client.aclose = AsyncMock()
    mock_cloud_response = MagicMock()
    mock_cloud_response.choices = [MagicMock(message=MagicMock(content="Cloud Response"))]
    mock_cloud_client.chat.completions.create = AsyncMock(return_value=mock_cloud_response)

    mock_openai.return_value = mock_cloud_client

    with patch.object(credential_service, "get_active_provider", AsyncMock(return_value=mock_provider_config)), \
         patch.object(credential_service, "get_credential", AsyncMock(return_value="0")), \
         patch.object(credential_service, "set_active_tier") as mock_set_tier, \
         patch.object(hybrid_router, "capability_matrix", mock_matrix):

        async with get_llm_client() as client:
            res = await client.chat.completions.create(messages=[{"role": "user", "content": "Hello"}])
            assert res.choices[0].message.content == "Cloud Response"
            mock_set_tier.assert_called_with(1)