46 lines
1.5 KiB
Python
46 lines
1.5 KiB
Python
import pytest
|
|
from unittest.mock import MagicMock, patch
|
|
from ea_chatbot.graph.nodes.planner import planner_node
|
|
|
|
@pytest.fixture
|
|
def mock_state():
|
|
return {
|
|
"messages": [],
|
|
"question": "Show me results for New Jersey",
|
|
"analysis": {
|
|
# "requires_dataset" removed as it's no longer used
|
|
"expert": "Data Analyst",
|
|
"data": "NJ data",
|
|
"unknown": "results",
|
|
"condition": "state=NJ"
|
|
},
|
|
"next_action": "plan",
|
|
"plan": None
|
|
}
|
|
|
|
@patch("ea_chatbot.graph.nodes.planner.get_llm_model")
|
|
@patch("ea_chatbot.utils.database_inspection.get_data_summary")
|
|
def test_planner_node(mock_get_summary, mock_get_llm, mock_state):
|
|
"""Test planner node with unified prompt."""
|
|
mock_get_summary.return_value = "Column: Name, Type: text"
|
|
|
|
mock_llm = MagicMock()
|
|
mock_get_llm.return_value = mock_llm
|
|
|
|
from ea_chatbot.schemas import TaskPlanResponse, TaskPlanContext
|
|
mock_plan = TaskPlanResponse(
|
|
goal="Get NJ results",
|
|
reflection="The user wants NJ results",
|
|
context=TaskPlanContext(initial_context="NJ data", assumptions=[], constraints=[]),
|
|
steps=["Step 1: Load data", "Step 2: Filter by NJ"]
|
|
)
|
|
mock_llm.with_structured_output.return_value.invoke.return_value = mock_plan
|
|
|
|
result = planner_node(mock_state)
|
|
|
|
assert "plan" in result
|
|
assert "Step 1: Load data" in result["plan"]
|
|
assert "Step 2: Filter by NJ" in result["plan"]
|
|
|
|
# Verify helper was called
|
|
mock_get_summary.assert_called_once() |