57 lines
1.5 KiB
Python
57 lines
1.5 KiB
Python
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"
|