Paul Bailey
f56f8086e5
- Updated RoleDetector to use OllamaService for AI-based speaker detection - Implemented _detect_with_ollama() method with availability check - Added fallback logic: tries Ollama first, then pattern matching - Implemented role merging and deduplication across both methods - Handles Ollama unavailable gracefully with logging - Logs detection method used (AI vs pattern) - Updated main.py endpoint to await async detect_roles - Added comprehensive integration tests Requirements: 14.1, 14.4, 14.5
176 lines
6.6 KiB
Python
176 lines
6.6 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Unit tests for Ollama integration in RoleDetector.
|
|
Tests the integration between RoleDetector and OllamaService.
|
|
"""
|
|
import pytest
|
|
import asyncio
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
from services.role_detector import RoleDetector
|
|
from services.ollama_service import OllamaService
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_detect_roles_with_ollama_available():
|
|
"""Test role detection when Ollama is available and returns speakers."""
|
|
# Create mock OllamaService
|
|
mock_ollama = AsyncMock(spec=OllamaService)
|
|
mock_ollama.is_available.return_value = True
|
|
mock_ollama.detect_speakers.return_value = ["Alice", "Bob", "Charlie"]
|
|
|
|
# Create RoleDetector with mock
|
|
detector = RoleDetector(ollama_service=mock_ollama)
|
|
|
|
# Test with text that has no pattern markers
|
|
text = "This is a conversation between Alice and Bob. Charlie joins later."
|
|
roles = await detector.detect_roles(text)
|
|
|
|
# Verify Ollama was called
|
|
mock_ollama.is_available.assert_called_once()
|
|
mock_ollama.detect_speakers.assert_called_once_with(text)
|
|
|
|
# Verify roles were detected
|
|
role_labels = [r['label'] for r in roles]
|
|
assert "Alice" in role_labels
|
|
assert "Bob" in role_labels
|
|
assert "Charlie" in role_labels
|
|
print("✓ PASS: Ollama detection works when available")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_detect_roles_with_ollama_unavailable():
|
|
"""Test role detection falls back to pattern matching when Ollama unavailable."""
|
|
# Create mock OllamaService that's unavailable
|
|
mock_ollama = AsyncMock(spec=OllamaService)
|
|
mock_ollama.is_available.return_value = False
|
|
|
|
# Create RoleDetector with mock
|
|
detector = RoleDetector(ollama_service=mock_ollama)
|
|
|
|
# Test with text that has pattern markers
|
|
text = "Alice: Hello there!\nBob: Hi Alice!"
|
|
roles = await detector.detect_roles(text)
|
|
|
|
# Verify Ollama availability was checked
|
|
mock_ollama.is_available.assert_called_once()
|
|
# Verify detect_speakers was NOT called since unavailable
|
|
mock_ollama.detect_speakers.assert_not_called()
|
|
|
|
# Verify pattern matching still worked
|
|
role_labels = [r['label'] for r in roles]
|
|
assert "Alice" in role_labels
|
|
assert "Bob" in role_labels
|
|
print("✓ PASS: Falls back to pattern matching when Ollama unavailable")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_detect_roles_with_ollama_error():
|
|
"""Test role detection handles Ollama errors gracefully."""
|
|
# Create mock OllamaService that throws error
|
|
mock_ollama = AsyncMock(spec=OllamaService)
|
|
mock_ollama.is_available.return_value = True
|
|
mock_ollama.detect_speakers.side_effect = Exception("Ollama connection failed")
|
|
|
|
# Create RoleDetector with mock
|
|
detector = RoleDetector(ollama_service=mock_ollama)
|
|
|
|
# Test with text that has pattern markers
|
|
text = "Alice: Hello there!\nBob: Hi Alice!"
|
|
roles = await detector.detect_roles(text)
|
|
|
|
# Verify Ollama was attempted
|
|
mock_ollama.is_available.assert_called_once()
|
|
mock_ollama.detect_speakers.assert_called_once()
|
|
|
|
# Verify pattern matching still worked as fallback
|
|
role_labels = [r['label'] for r in roles]
|
|
assert "Alice" in role_labels
|
|
assert "Bob" in role_labels
|
|
print("✓ PASS: Handles Ollama errors gracefully and falls back to patterns")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_detect_roles_merges_ollama_and_pattern_results():
|
|
"""Test that roles from both Ollama and pattern matching are merged."""
|
|
# Create mock OllamaService
|
|
mock_ollama = AsyncMock(spec=OllamaService)
|
|
mock_ollama.is_available.return_value = True
|
|
mock_ollama.detect_speakers.return_value = ["Alice", "Charlie"]
|
|
|
|
# Create RoleDetector with mock
|
|
detector = RoleDetector(ollama_service=mock_ollama)
|
|
|
|
# Test with text that has pattern markers for different roles
|
|
text = "Alice: Hello!\nBob: Hi there!\nCharlie appears later."
|
|
roles = await detector.detect_roles(text)
|
|
|
|
# Verify both methods were used
|
|
mock_ollama.detect_speakers.assert_called_once()
|
|
|
|
# Verify roles from both sources are present
|
|
role_labels = [r['label'] for r in roles]
|
|
assert "Alice" in role_labels # From both Ollama and pattern
|
|
assert "Bob" in role_labels # From pattern only
|
|
assert "Charlie" in role_labels # From Ollama only
|
|
print("✓ PASS: Merges results from both Ollama and pattern matching")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_detect_roles_deduplicates_across_methods():
|
|
"""Test that duplicate roles from Ollama and patterns are deduplicated."""
|
|
# Create mock OllamaService
|
|
mock_ollama = AsyncMock(spec=OllamaService)
|
|
mock_ollama.is_available.return_value = True
|
|
mock_ollama.detect_speakers.return_value = ["Alice", "Bob"]
|
|
|
|
# Create RoleDetector with mock
|
|
detector = RoleDetector(ollama_service=mock_ollama)
|
|
|
|
# Test with text where patterns also detect Alice and Bob
|
|
text = "Alice: Hello!\nBob: Hi there!"
|
|
roles = await detector.detect_roles(text)
|
|
|
|
# Verify roles are deduplicated
|
|
role_labels = [r['label'] for r in roles]
|
|
assert role_labels.count("Alice") == 1
|
|
assert role_labels.count("Bob") == 1
|
|
assert len(role_labels) == 2
|
|
print("✓ PASS: Deduplicates roles across both detection methods")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_detect_roles_returns_narrator_when_no_roles_found():
|
|
"""Test that default Narrator is returned when no roles detected by either method."""
|
|
# Create mock OllamaService that returns empty
|
|
mock_ollama = AsyncMock(spec=OllamaService)
|
|
mock_ollama.is_available.return_value = True
|
|
mock_ollama.detect_speakers.return_value = []
|
|
|
|
# Create RoleDetector with mock
|
|
detector = RoleDetector(ollama_service=mock_ollama)
|
|
|
|
# Test with text that has no role markers
|
|
text = "This is just plain text without any speakers."
|
|
roles = await detector.detect_roles(text)
|
|
|
|
# Verify default Narrator is returned
|
|
assert len(roles) == 1
|
|
assert roles[0]['label'] == "Narrator"
|
|
print("✓ PASS: Returns default Narrator when no roles found")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
# Run tests
|
|
print("Testing Ollama Integration in RoleDetector...\n")
|
|
|
|
asyncio.run(test_detect_roles_with_ollama_available())
|
|
asyncio.run(test_detect_roles_with_ollama_unavailable())
|
|
asyncio.run(test_detect_roles_with_ollama_error())
|
|
asyncio.run(test_detect_roles_merges_ollama_and_pattern_results())
|
|
asyncio.run(test_detect_roles_deduplicates_across_methods())
|
|
asyncio.run(test_detect_roles_returns_narrator_when_no_roles_found())
|
|
|
|
print("\n" + "="*50)
|
|
print("All integration tests passed!")
|
|
print("="*50)
|