tests/test_router.py

from unittest.mock import patch, MagicMock
from rag.router import classify_query

def test_classify_query_exact_policy():
    """Verify that exact classification response 'policy' is correctly resolved."""
    with patch('rag.router.groq_client') as mock_client:
        mock_response = MagicMock()
        mock_response.choices = [MagicMock()]
        mock_response.choices[0].message.content = "policy"
        mock_client.chat.completions.create.return_value = mock_response
        
        category = classify_query("What is the return policy?")
        assert category == "policy"
        mock_client.chat.completions.create.assert_called_once()

def test_classify_query_exact_product():
    """Verify that exact classification response 'product' is correctly resolved."""
    with patch('rag.router.groq_client') as mock_client:
        mock_response = MagicMock()
        mock_response.choices = [MagicMock()]
        mock_response.choices[0].message.content = "product"
        mock_client.chat.completions.create.return_value = mock_response
        
        category = classify_query("Do you have toy medieval castle?")
        assert category == "product"

def test_classify_query_exact_none():
    """Verify that exact classification response 'none' is correctly resolved."""
    with patch('rag.router.groq_client') as mock_client:
        mock_response = MagicMock()
        mock_response.choices = [MagicMock()]
        mock_response.choices[0].message.content = "none"
        mock_client.chat.completions.create.return_value = mock_response
        
        category = classify_query("hello there")
        assert category == "none"

def test_classify_query_partial_match():
    """Verify that partial match fallbacks work (e.g. LLM responds with a phrase containing the keyword)."""
    with patch('rag.router.groq_client') as mock_client:
        mock_response = MagicMock()
        mock_response.choices = [MagicMock()]
        mock_response.choices[0].message.content = "This is a product classification"
        mock_client.chat.completions.create.return_value = mock_response
        
        category = classify_query("toy knights set")
        assert category == "product"

def test_classify_query_fallback_on_unrecognized():
    """Verify that unrecognized classifications default to 'none'."""
    with patch('rag.router.groq_client') as mock_client:
        mock_response = MagicMock()
        mock_response.choices = [MagicMock()]
        mock_response.choices[0].message.content = "something completely unrelated"
        mock_client.chat.completions.create.return_value = mock_response
        
        category = classify_query("is there a sale?")
        assert category == "none"

def test_classify_query_exception_fallback():
    """Verify that if the LLM request raises an exception, the router falls back to 'none' gracefully."""
    with patch('rag.router.groq_client') as mock_client:
        mock_client.chat.completions.create.side_effect = Exception("API connection failure")
        
        category = classify_query("any question")
        assert category == "none"