voice-assistant/tests/test_voice_assistant.py

"""
Unit tests for VoiceAssistant handler.
"""
import base64
import pytest
from unittest.mock import AsyncMock, MagicMock, patch

# Import after environment is set up in conftest
from voice_assistant import VoiceAssistant, VoiceSettings


class TestVoiceSettings:
    """Tests for VoiceSettings configuration."""
    
    def test_default_settings(self):
        """Test default settings values."""
        settings = VoiceSettings()
        
        assert settings.service_name == "voice-assistant"
        assert settings.rag_top_k == 10
        assert settings.rag_rerank_top_k == 5
        assert settings.rag_collection == "documents"
        assert settings.stt_language is None  # Auto-detect
        assert settings.tts_language == "en"
        assert settings.include_transcription is True
        assert settings.include_sources is False
    
    def test_custom_settings(self, monkeypatch):
        """Test settings from environment."""
        monkeypatch.setenv("RAG_TOP_K", "20")
        monkeypatch.setenv("RAG_COLLECTION", "custom_collection")
        
        # Note: Would need to re-instantiate settings to pick up env vars
        settings = VoiceSettings(
            rag_top_k=20,
            rag_collection="custom_collection"
        )
        
        assert settings.rag_top_k == 20
        assert settings.rag_collection == "custom_collection"


class TestVoiceAssistant:
    """Tests for VoiceAssistant handler."""
    
    @pytest.fixture
    def handler(self):
        """Create handler with mocked clients."""
        with patch("voice_assistant.STTClient"), \
             patch("voice_assistant.EmbeddingsClient"), \
             patch("voice_assistant.RerankerClient"), \
             patch("voice_assistant.LLMClient"), \
             patch("voice_assistant.TTSClient"), \
             patch("voice_assistant.MilvusClient"):
            
            handler = VoiceAssistant()
            
            # Setup mock clients
            handler.stt = AsyncMock()
            handler.embeddings = AsyncMock()
            handler.reranker = AsyncMock()
            handler.llm = AsyncMock()
            handler.tts = AsyncMock()
            handler.milvus = AsyncMock()
            handler.nats = AsyncMock()
            
            yield handler
    
    def test_init(self, handler):
        """Test handler initialization."""
        assert handler.subject == "voice.request"
        assert handler.queue_group == "voice-assistants"
        assert handler.voice_settings.service_name == "voice-assistant"
    
    @pytest.mark.asyncio
    async def test_handle_message_success(
        self,
        handler,
        mock_nats_message,
        mock_voice_request,
        sample_embedding,
        sample_documents,
        sample_reranked,
    ):
        """Test successful voice request handling."""
        # Setup mocks
        handler.stt.transcribe.return_value = {"text": "What is machine learning?"}
        handler.embeddings.embed_single.return_value = sample_embedding
        handler.milvus.search_with_texts.return_value = sample_documents
        handler.reranker.rerank.return_value = sample_reranked
        handler.llm.generate.return_value = "Machine learning is a type of AI."
        handler.tts.synthesize.return_value = b"audio_bytes"
        
        # Execute
        result = await handler.handle_message(mock_nats_message, mock_voice_request)
        
        # Verify
        assert result["request_id"] == "test-request-123"
        assert result["response"] == "Machine learning is a type of AI."
        assert "audio" in result
        assert result["transcription"] == "What is machine learning?"
        
        # Verify pipeline was called
        handler.stt.transcribe.assert_called_once()
        handler.embeddings.embed_single.assert_called_once()
        handler.milvus.search_with_texts.assert_called_once()
        handler.reranker.rerank.assert_called_once()
        handler.llm.generate.assert_called_once()
        handler.tts.synthesize.assert_called_once()
    
    @pytest.mark.asyncio
    async def test_handle_message_empty_transcription(
        self,
        handler,
        mock_nats_message,
        mock_voice_request,
    ):
        """Test handling when transcription is empty."""
        handler.stt.transcribe.return_value = {"text": ""}
        
        result = await handler.handle_message(mock_nats_message, mock_voice_request)
        
        assert "error" in result
        assert result["error"] == "Could not transcribe audio"
        
        # Verify pipeline stopped after transcription
        handler.embeddings.embed_single.assert_not_called()
    
    @pytest.mark.asyncio
    async def test_handle_message_with_sources(
        self,
        handler,
        mock_nats_message,
        mock_voice_request,
        sample_embedding,
        sample_documents,
        sample_reranked,
    ):
        """Test response includes sources when enabled."""
        handler.voice_settings.include_sources = True
        
        # Setup mocks
        handler.stt.transcribe.return_value = {"text": "Hello"}
        handler.embeddings.embed_single.return_value = sample_embedding
        handler.milvus.search_with_texts.return_value = sample_documents
        handler.reranker.rerank.return_value = sample_reranked
        handler.llm.generate.return_value = "Hi there!"
        handler.tts.synthesize.return_value = b"audio"
        
        result = await handler.handle_message(mock_nats_message, mock_voice_request)
        
        assert "sources" in result
        assert len(result["sources"]) <= 3
    
    def test_build_context(self, handler):
        """Test context building from documents."""
        documents = [
            {"document": "First doc content"},
            {"document": "Second doc content"},
        ]
        
        context = handler._build_context(documents)
        
        assert "First doc content" in context
        assert "Second doc content" in context
    
    @pytest.mark.asyncio
    async def test_setup_initializes_clients(self):
        """Test that setup initializes all clients."""
        with patch("voice_assistant.STTClient") as stt_cls, \
             patch("voice_assistant.EmbeddingsClient") as emb_cls, \
             patch("voice_assistant.RerankerClient") as rer_cls, \
             patch("voice_assistant.LLMClient") as llm_cls, \
             patch("voice_assistant.TTSClient") as tts_cls, \
             patch("voice_assistant.MilvusClient") as mil_cls:
            
            mil_cls.return_value.connect = AsyncMock()
            
            handler = VoiceAssistant()
            await handler.setup()
            
            stt_cls.assert_called_once()
            emb_cls.assert_called_once()
            rer_cls.assert_called_once()
            llm_cls.assert_called_once()
            tts_cls.assert_called_once()
            mil_cls.assert_called_once()
    
    @pytest.mark.asyncio
    async def test_teardown_closes_clients(self, handler):
        """Test that teardown closes all clients."""
        await handler.teardown()
        
        handler.stt.close.assert_called_once()
        handler.embeddings.close.assert_called_once()
        handler.reranker.close.assert_called_once()
        handler.llm.close.assert_called_once()
        handler.tts.close.assert_called_once()
        handler.milvus.close.assert_called_once()