fix(orchestrator): Apply refinements from code review

This commit is contained in:
Yunxiao Xu
2026-02-23 15:46:21 -08:00
parent c5cf4b38a1
commit 2cfbc5d1d0
19 changed files with 252 additions and 33 deletions

View File

@@ -0,0 +1,42 @@
import pytest
from ea_chatbot.graph.workflow import create_workflow, data_analyst_worker_node
from ea_chatbot.graph.state import AgentState
from unittest.mock import MagicMock, patch
from langchain_core.messages import HumanMessage, AIMessage
def test_worker_merge_sets_summary_for_reflector():
"""Verify that worker node (wrapper) sets the 'summary' field for the Reflector."""
state = AgentState(
messages=[HumanMessage(content="test")],
question="test",
checklist=[{"task": "Analyze data", "worker": "data_analyst"}],
current_step=0,
iterations=0,
vfs={},
plots=[],
dfs={},
next_action="",
analysis={},
summary="Initial Planner Summary" # Stale summary
)
# Mock the compiled worker subgraph to return a specific result
with patch("ea_chatbot.graph.workflow._DATA_ANALYST_WORKER") as mock_worker:
mock_worker.invoke.return_value = {
"result": "Actual Worker Findings",
"messages": [AIMessage(content="Internal")],
"vfs_state": {},
"plots": []
}
# Execute the wrapper node
updates = data_analyst_worker_node(state)
# Verify that 'summary' is in updates and has the worker result
assert "summary" in updates
assert updates["summary"] == "Actual Worker Findings"
# When applied to state, it should overwrite the stale summary
state.update(updates)
assert state["summary"] == "Actual Worker Findings"