audiobook/backend/test_ollama_service.py

"""
Tests for Ollama integration service.
"""
import pytest
import json
from unittest.mock import AsyncMock, patch, MagicMock
from services.ollama_service import OllamaService


@pytest.mark.asyncio
async def test_ollama_service_initialization():
    """Test OllamaService initializes with correct defaults."""
    service = OllamaService()
    assert service.base_url == "http://localhost:11434"
    assert service.model == "llama3:8b"
    assert service.timeout == 30.0


@pytest.mark.asyncio
async def test_ollama_service_custom_base_url():
    """Test OllamaService accepts custom base URL."""
    custom_url = "http://custom-ollama:11434"
    service = OllamaService(base_url=custom_url)
    assert service.base_url == custom_url


@pytest.mark.asyncio
async def test_build_prompt():
    """Test prompt builder creates correct format."""
    service = OllamaService()
    text = "Alice: Hello there!\nBob: Hi Alice!"
    prompt = service._build_prompt(text)
    
    assert "speakers" in prompt
    assert "JSON" in prompt
    assert text in prompt or text[:2000] in prompt


@pytest.mark.asyncio
async def test_parse_response_with_speakers_key():
    """Test parsing JSON response with 'speakers' key."""
    service = OllamaService()
    response = '{"speakers": ["Alice", "Bob", "Narrator"]}'
    speakers = service._parse_response(response)
    
    assert speakers == ["Alice", "Bob", "Narrator"]


@pytest.mark.asyncio
async def test_parse_response_with_characters_key():
    """Test parsing JSON response with 'characters' key."""
    service = OllamaService()
    response = '{"characters": ["Alice", "Bob"]}'
    speakers = service._parse_response(response)
    
    assert speakers == ["Alice", "Bob"]


@pytest.mark.asyncio
async def test_parse_response_with_roles_key():
    """Test parsing JSON response with 'roles' key."""
    service = OllamaService()
    response = '{"roles": ["Narrator", "Alice"]}'
    speakers = service._parse_response(response)
    
    assert speakers == ["Narrator", "Alice"]


@pytest.mark.asyncio
async def test_parse_response_with_array():
    """Test parsing JSON response that is a direct array."""
    service = OllamaService()
    response = '["Alice", "Bob"]'
    speakers = service._parse_response(response)
    
    assert speakers == ["Alice", "Bob"]


@pytest.mark.asyncio
async def test_parse_response_empty():
    """Test parsing empty response."""
    service = OllamaService()
    speakers = service._parse_response("")
    assert speakers == []


@pytest.mark.asyncio
async def test_parse_response_invalid_json():
    """Test parsing invalid JSON returns empty list."""
    service = OllamaService()
    response = "This is not JSON"
    speakers = service._parse_response(response)
    assert speakers == []


@pytest.mark.asyncio
async def test_detect_speakers_success():
    """Test successful speaker detection."""
    service = OllamaService()
    text = "Alice: Hello!\nBob: Hi there!"
    
    mock_response = MagicMock()
    mock_response.json.return_value = {
        "response": '{"speakers": ["Alice", "Bob"]}'
    }
    mock_response.raise_for_status = MagicMock()
    
    with patch('httpx.AsyncClient') as mock_client:
        mock_client.return_value.__aenter__.return_value.post = AsyncMock(return_value=mock_response)
        
        speakers = await service.detect_speakers(text)
        assert speakers == ["Alice", "Bob"]


@pytest.mark.asyncio
async def test_detect_speakers_timeout():
    """Test speaker detection handles timeout."""
    service = OllamaService()
    text = "Some text"
    
    with patch('httpx.AsyncClient') as mock_client:
        import httpx
        mock_client.return_value.__aenter__.return_value.post = AsyncMock(
            side_effect=httpx.TimeoutException("Timeout")
        )
        
        with pytest.raises(Exception, match="timed out"):
            await service.detect_speakers(text)


@pytest.mark.asyncio
async def test_detect_speakers_http_error():
    """Test speaker detection handles HTTP errors."""
    service = OllamaService()
    text = "Some text"
    
    with patch('httpx.AsyncClient') as mock_client:
        import httpx
        mock_client.return_value.__aenter__.return_value.post = AsyncMock(
            side_effect=httpx.HTTPError("Connection failed")
        )
        
        with pytest.raises(Exception, match="Ollama service error"):
            await service.detect_speakers(text)


@pytest.mark.asyncio
async def test_is_available_success():
    """Test is_available returns True when service responds."""
    service = OllamaService()
    
    mock_response = MagicMock()
    mock_response.status_code = 200
    
    with patch('httpx.AsyncClient') as mock_client:
        mock_client.return_value.__aenter__.return_value.get = AsyncMock(return_value=mock_response)
        
        available = await service.is_available()
        assert available is True


@pytest.mark.asyncio
async def test_is_available_failure():
    """Test is_available returns False when service is down."""
    service = OllamaService()
    
    with patch('httpx.AsyncClient') as mock_client:
        import httpx
        mock_client.return_value.__aenter__.return_value.get = AsyncMock(
            side_effect=httpx.ConnectError("Connection refused")
        )
        
        available = await service.is_available()
        assert available is False