46 lines
1.5 KiB
Python
46 lines
1.5 KiB
Python
import pytest
|
|
import logging
|
|
from unittest.mock import MagicMock, patch
|
|
from ea_chatbot.graph.nodes.query_analyzer import query_analyzer_node, QueryAnalysis
|
|
|
|
@pytest.fixture
|
|
def mock_state():
|
|
return {
|
|
"messages": [],
|
|
"question": "Show me the 2024 results for Florida",
|
|
"analysis": None,
|
|
"next_action": ""
|
|
}
|
|
|
|
@patch("ea_chatbot.graph.nodes.query_analyzer.get_llm_model")
|
|
@patch("ea_chatbot.graph.nodes.query_analyzer.get_logger")
|
|
def test_query_analyzer_logs_actions(mock_get_logger, mock_get_llm, mock_state):
|
|
"""Test that query_analyzer_node logs its main actions."""
|
|
# Mock Logger
|
|
mock_logger = MagicMock()
|
|
mock_get_logger.return_value = mock_logger
|
|
|
|
# Mock LLM
|
|
mock_llm_instance = MagicMock()
|
|
mock_get_llm.return_value = mock_llm_instance
|
|
mock_structured_llm = MagicMock()
|
|
mock_llm_instance.with_structured_output.return_value = mock_structured_llm
|
|
|
|
expected_analysis = QueryAnalysis(
|
|
data_required=["2024 results", "Florida"],
|
|
unknowns=[],
|
|
ambiguities=[],
|
|
conditions=[],
|
|
next_action="plan"
|
|
)
|
|
mock_structured_llm.invoke.return_value = expected_analysis
|
|
|
|
query_analyzer_node(mock_state)
|
|
|
|
# Check that logger was called
|
|
# We expect at least one log at the start and one at the end
|
|
assert mock_logger.info.called
|
|
|
|
# Verify specific log messages if we decide on them
|
|
# For now, just ensuring it's called is enough for Red phase
|