tests/test_chat_service.py

from unittest.mock import patch, MagicMock
from services.chat_service import (
    resolve_category,
    get_retrieved_context,
    build_llm_messages,
    generate_llm_response,
    chat
)

class MockDocument:
    def __init__(self, page_content, metadata=None):
        self.page_content = page_content
        self.metadata = metadata or {}

# 1. Tests for resolve_category
def test_resolve_category_manual():
    """Verify that manual doc_type overrides map directly to their corresponding categories."""
    assert resolve_category("hello", "Policy") == "policy"
    assert resolve_category("hello", "Product") == "product"
    assert resolve_category("hello", "No Filter") == "none"

@patch("services.chat_service.classify_query")
def test_resolve_category_auto(mock_classify):
    """Verify that 'Auto (AI Router)' delegates to the classifier."""
    mock_classify.return_value = "policy"
    assert resolve_category("return policy?", "Auto (AI Router)") == "policy"
    mock_classify.assert_called_once_with("return policy?")

# 2. Tests for get_retrieved_context
def test_get_retrieved_context_none():
    """Verify that 'none' category skips context retrieval entirely."""
    with patch("services.chat_service.retrieve_all") as mock_retrieve:
        context = get_retrieved_context("hi", "none")
        assert context == ""
        mock_retrieve.assert_not_called()

@patch("services.chat_service.retrieve_all")
def test_get_retrieved_context_active(mock_retrieve):
    """Verify that active categories retrieve and format documents."""
    mock_retrieve.return_value = [
        MockDocument("Doc A content"),
        MockDocument("Doc B content")
    ]
    context = get_retrieved_context("question", "policy")
    assert context == "Doc A content\n\nDoc B content"
    mock_retrieve.assert_called_once_with("question", "policy")

# 3. Tests for build_llm_messages
def test_build_llm_messages_no_context():
    """Verify messages layout when context is empty."""
    messages = build_llm_messages("my question", "")
    # Should have system prompt and user question
    assert len(messages) == 2
    assert messages[0]["role"] == "system"
    assert messages[1]["role"] == "user"
    assert messages[1]["content"] == "my question"

def test_build_llm_messages_with_context():
    """Verify messages layout when context is provided."""
    messages = build_llm_messages("my question", "some retrieved facts")
    # Should have system prompt, context injection prompt, and user question
    assert len(messages) == 3
    assert messages[0]["role"] == "system"
    assert messages[1]["role"] == "system"
    assert "some retrieved facts" in messages[1]["content"]
    assert messages[2]["role"] == "user"
    assert messages[2]["content"] == "my question"

# 4. Tests for generate_llm_response
@patch("services.chat_service.groq_client")
def test_generate_llm_response(mock_client):
    """Verify response parsing and configuration settings for Groq LLM API."""
    mock_response = MagicMock()
    mock_response.choices = [MagicMock()]
    mock_response.choices[0].message.content = "Excellent toy choice!"
    mock_client.chat.completions.create.return_value = mock_response

    messages = [{"role": "user", "content": "hi"}]
    answer = generate_llm_response(messages)

    assert answer == "Excellent toy choice!"
    mock_client.chat.completions.create.assert_called_once_with(
        model="llama-3.3-70b-versatile",
        messages=messages,
        max_tokens=400,
        temperature=0.6
    )

# 5. Coordinating flow chat workflow tests
@patch("services.chat_service.resolve_category")
@patch("services.chat_service.get_retrieved_context")
@patch("services.chat_service.build_llm_messages")
@patch("services.chat_service.generate_llm_response")
def test_chat_success(mock_generate, mock_build, mock_get, mock_resolve):
    """Verify the coordinated end-to-end chat routing workflow works on success."""
    mock_resolve.return_value = "product"
    mock_get.return_value = "Toy facts"
    mock_build.return_value = [{"role": "user", "content": "Query"}]
    mock_generate.return_value = "LLM answer"

    response = chat("What is toy X?", "Auto (AI Router)")

    assert response == "LLM answer"
    mock_resolve.assert_called_once_with("What is toy X?", "Auto (AI Router)")
    mock_get.assert_called_once_with("What is toy X?", "product")
    mock_build.assert_called_once_with("What is toy X?", "Toy facts", None)
    mock_generate.assert_called_once()

def test_chat_exception_handling():
    """Verify that exceptions are caught and return a safe error string."""
    with patch("services.chat_service.resolve_category", side_effect=Exception("Database crash")):
        response = chat("any question", "Auto (AI Router)")
        assert "An error occurred while formulating a response:" in response
        assert "Database crash" in response

def test_build_llm_messages_with_history():
    """Verify message context structure layout when conversation history is provided."""
    history = [
        {"role": "user", "content": "hi Maya"},
        {"role": "assistant", "content": "Hello! How can I help you?"}
    ]
    messages = build_llm_messages("my second question", "some retrieved facts", history)
    
    # Standard format: system prompt, context prompt, history user, history assistant, current user question
    assert len(messages) == 5
    assert messages[0]["role"] == "system"
    assert messages[1]["role"] == "system"
    assert "some retrieved facts" in messages[1]["content"]
    assert messages[2]["role"] == "user"
    assert messages[2]["content"] == "hi Maya"
    assert messages[3]["role"] == "assistant"
    assert messages[3]["content"] == "Hello! How can I help you?"
    assert messages[4]["role"] == "user"
    assert messages[4]["content"] == "my second question"

@patch("services.chat_service.resolve_category")
@patch("services.chat_service.get_retrieved_context")
@patch("services.chat_service.build_llm_messages")
@patch("services.chat_service.generate_llm_response")
def test_chat_with_history_success(mock_generate, mock_build, mock_get, mock_resolve):
    """Verify the coordinated end-to-end chat workflow with history works on success."""
    mock_resolve.return_value = "product"
    mock_get.return_value = "Toy facts"
    mock_build.return_value = [{"role": "user", "content": "Query"}]
    mock_generate.return_value = "LLM answer"
    
    history = [{"role": "user", "content": "hi"}]
    response = chat("What is toy X?", "Auto (AI Router)", history)
    
    assert response == "LLM answer"
    mock_resolve.assert_called_once_with("What is toy X?", "Auto (AI Router)")
    mock_get.assert_called_once_with("What is toy X?", "product")
    mock_build.assert_called_once_with("What is toy X?", "Toy facts", history)
    mock_generate.assert_called_once()