87 lines
2.5 KiB
Python
87 lines
2.5 KiB
Python
import pytest
|
|
from unittest.mock import MagicMock
|
|
from ea_chatbot.graph.workflow import create_workflow
|
|
from ea_chatbot.graph.state import AgentState
|
|
from langchain_core.messages import AIMessage, HumanMessage
|
|
|
|
def test_orchestrator_full_flow():
|
|
"""Verify the full Orchestrator-Workers flow via direct node injection."""
|
|
|
|
mock_analyzer = MagicMock()
|
|
mock_planner = MagicMock()
|
|
mock_delegate = MagicMock()
|
|
mock_worker = MagicMock()
|
|
mock_reflector = MagicMock()
|
|
mock_synthesizer = MagicMock()
|
|
mock_summarize_conv = MagicMock()
|
|
|
|
# 1. Analyzer: Proceed to planning
|
|
mock_analyzer.return_value = {"next_action": "plan"}
|
|
|
|
# 2. Planner: Generate checklist
|
|
mock_planner.return_value = {
|
|
"checklist": [{"task": "T1", "worker": "data_analyst"}],
|
|
"current_step": 0
|
|
}
|
|
|
|
# 3. Delegate: Route to data_analyst
|
|
mock_delegate.side_effect = [
|
|
{"next_action": "data_analyst"}, # First call
|
|
{"next_action": "summarize"} # Second call (after reflector)
|
|
]
|
|
|
|
# 4. Worker: Success
|
|
mock_worker.return_value = {
|
|
"messages": [AIMessage(content="Worker result")],
|
|
"vfs": {"res.txt": "data"}
|
|
}
|
|
|
|
# 5. Reflector: Advance
|
|
mock_reflector.return_value = {
|
|
"current_step": 1,
|
|
"next_action": "delegate"
|
|
}
|
|
|
|
# 6. Synthesizer: Final answer
|
|
mock_synthesizer.return_value = {
|
|
"messages": [AIMessage(content="Final synthesized answer")],
|
|
"next_action": "end"
|
|
}
|
|
|
|
# 7. Summarize Conv: End
|
|
mock_summarize_conv.return_value = {"summary": "Done"}
|
|
|
|
# Create workflow with injected mocks
|
|
app = create_workflow(
|
|
query_analyzer=mock_analyzer,
|
|
planner=mock_planner,
|
|
delegate=mock_delegate,
|
|
data_analyst_worker=mock_worker,
|
|
reflector=mock_reflector,
|
|
synthesizer=mock_synthesizer,
|
|
summarize_conversation=mock_summarize_conv
|
|
)
|
|
|
|
initial_state = AgentState(
|
|
messages=[HumanMessage(content="Explain results")],
|
|
question="Explain results",
|
|
analysis={},
|
|
next_action="",
|
|
iterations=0,
|
|
checklist=[],
|
|
current_step=0,
|
|
vfs={},
|
|
plots=[],
|
|
dfs={}
|
|
)
|
|
|
|
final_state = app.invoke(initial_state)
|
|
|
|
assert mock_analyzer.called
|
|
assert mock_planner.called
|
|
assert mock_delegate.call_count == 2
|
|
assert mock_worker.called
|
|
assert mock_reflector.called
|
|
assert mock_synthesizer.called
|
|
assert "Final synthesized answer" in [m.content for m in final_state["messages"]]
|