test(orchestrator): Add integration tests for the Orchestrator-Workers loop
This commit is contained in:
86
backend/tests/test_orchestrator_loop.py
Normal file
86
backend/tests/test_orchestrator_loop.py
Normal file
@@ -0,0 +1,86 @@
|
||||
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"]]
|
||||
Reference in New Issue
Block a user