from ea_chatbot.graph.nodes.delegate import delegate_node
from ea_chatbot.graph.state import AgentState

def test_delegate_node_data_analyst():
    """Verify that the delegate node routes to data_analyst."""
    state = AgentState(
        checklist=[{"task": "Analyze data", "worker": "data_analyst"}],
        current_step=0,
        messages=[],
        question="test",
        analysis={},
        next_action="",
        iterations=0,
        vfs={},
        plots=[],
        dfs={}
    )
    
    result = delegate_node(state)
    assert result["next_action"] == "data_analyst"

def test_delegate_node_researcher():
    """Verify that the delegate node routes to researcher."""
    state = AgentState(
        checklist=[{"task": "Search web", "worker": "researcher"}],
        current_step=0,
        messages=[],
        question="test",
        analysis={},
        next_action="",
        iterations=0,
        vfs={},
        plots=[],
        dfs={}
    )
    
    result = delegate_node(state)
    assert result["next_action"] == "researcher"

def test_delegate_node_finished():
    """Verify that the delegate node routes to summarize if checklist is complete."""
    state = AgentState(
        checklist=[{"task": "Task 1", "worker": "data_analyst"}],
        current_step=1, # Already finished the only task
        messages=[],
        question="test",
        analysis={},
        next_action="",
        iterations=0,
        vfs={},
        plots=[],
        dfs={}
    )
    
    result = delegate_node(state)
    assert result["next_action"] == "summarize"