Paul Bailey
73dc5695ca
- Sample up to 4 chunks of 5000 characters from different text positions - Distribute samples evenly across beginning, middle, and end of text - Deduplicate speakers across all samples for comprehensive detection - Log individual sample results for debugging - Handle short texts efficiently (single sample) - Updated tests to verify multi-sample behavior This ensures speakers introduced anywhere in long texts are detected, not just those appearing in the first 2000 characters.
238 lines
7.7 KiB
Python
238 lines
7.7 KiB
Python
"""
|
|
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_sample_text_short():
|
|
"""Test sampling short text returns full text."""
|
|
service = OllamaService()
|
|
text = "Short text"
|
|
samples = service._sample_text(text, chunk_size=5000, num_samples=4)
|
|
|
|
assert len(samples) == 1
|
|
assert samples[0] == text
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sample_text_long():
|
|
"""Test sampling long text returns multiple samples."""
|
|
service = OllamaService()
|
|
# Create text with position markers so we can verify different samples
|
|
text = "".join([f"Position{i:05d} " for i in range(2500)]) # ~25k characters
|
|
samples = service._sample_text(text, chunk_size=5000, num_samples=4)
|
|
|
|
assert len(samples) == 4
|
|
assert all(len(s) <= 5000 for s in samples)
|
|
# Samples should be from different parts - check they contain different position markers
|
|
assert "Position00000" in samples[0] # First sample should have early positions
|
|
assert "Position00000" not in samples[-1] # Last sample should have later positions
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sample_text_medium():
|
|
"""Test sampling medium text returns appropriate number of samples."""
|
|
service = OllamaService()
|
|
text = "B" * 8000 # 8k characters
|
|
samples = service._sample_text(text, chunk_size=5000, num_samples=4)
|
|
|
|
# Should get 2 samples (8000 / 5000 = 1.6, so 2 samples)
|
|
assert len(samples) >= 1
|
|
assert all(len(s) <= 5000 for s in samples)
|
|
|
|
|
|
@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
|
|
|
|
|
|
@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 with short text."""
|
|
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 set(speakers) == {"Alice", "Bob"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_detect_speakers_multiple_samples():
|
|
"""Test speaker detection with long text requiring multiple samples."""
|
|
service = OllamaService()
|
|
text = "A" * 25000 # Long text requiring multiple samples
|
|
|
|
# Mock responses for different samples
|
|
responses = [
|
|
MagicMock(json=lambda: {"response": '{"speakers": ["Alice", "Bob"]}'}),
|
|
MagicMock(json=lambda: {"response": '{"speakers": ["Charlie", "Bob"]}'}),
|
|
MagicMock(json=lambda: {"response": '{"speakers": ["Alice", "David"]}'}),
|
|
MagicMock(json=lambda: {"response": '{"speakers": ["Eve"]}'}),
|
|
]
|
|
|
|
for r in responses:
|
|
r.raise_for_status = MagicMock()
|
|
|
|
with patch('httpx.AsyncClient') as mock_client:
|
|
mock_client.return_value.__aenter__.return_value.post = AsyncMock(side_effect=responses)
|
|
|
|
speakers = await service.detect_speakers(text)
|
|
# Should deduplicate across all samples
|
|
assert set(speakers) == {"Alice", "Bob", "Charlie", "David", "Eve"}
|
|
|
|
|
|
@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
|