49 lines
1.6 KiB
Python
49 lines
1.6 KiB
Python
import pytest
|
|
from unittest.mock import MagicMock, patch
|
|
from langchain_core.messages import AIMessage
|
|
from ea_chatbot.graph.nodes.summarizer import summarizer_node
|
|
from ea_chatbot.graph.state import AgentState
|
|
|
|
@pytest.fixture
|
|
def mock_llm():
|
|
with patch("ea_chatbot.graph.nodes.summarizer.get_llm_model") as mock_get_llm:
|
|
mock_llm_instance = MagicMock()
|
|
mock_get_llm.return_value = mock_llm_instance
|
|
yield mock_llm_instance
|
|
|
|
def test_summarizer_node_success(mock_llm):
|
|
"""Test that summarizer_node invokes LLM with correct inputs and returns messages."""
|
|
state = {
|
|
"question": "What is the total count?",
|
|
"plan": "1. Run query\n2. Sum results",
|
|
"code_output": "The total is 100",
|
|
"messages": []
|
|
}
|
|
|
|
mock_llm.invoke.return_value = AIMessage(content="The final answer is 100.")
|
|
|
|
result = summarizer_node(state)
|
|
|
|
# Verify LLM was called
|
|
assert mock_llm.invoke.called
|
|
|
|
# Verify result structure
|
|
assert "messages" in result
|
|
assert len(result["messages"]) == 1
|
|
assert isinstance(result["messages"][0], AIMessage)
|
|
assert result["messages"][0].content == "The final answer is 100."
|
|
|
|
def test_summarizer_node_empty_state(mock_llm):
|
|
"""Test handling of empty or minimal state."""
|
|
state = {
|
|
"question": "Empty?",
|
|
"messages": []
|
|
}
|
|
|
|
mock_llm.invoke.return_value = AIMessage(content="No data provided.")
|
|
|
|
result = summarizer_node(state)
|
|
|
|
assert "messages" in result
|
|
assert result["messages"][0].content == "No data provided."
|