49 lines
1.8 KiB
Python
49 lines
1.8 KiB
Python
import pytest
|
|
from ea_chatbot.graph.workflow import create_workflow, data_analyst_worker_runnable
|
|
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(monkeypatch):
|
|
"""Verify that worker node (runnable) 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
|
|
)
|
|
|
|
# Create a mock for the invoke method
|
|
mock_invoke = MagicMock()
|
|
mock_invoke.return_value = {
|
|
"summary": "Actual Worker Findings",
|
|
"messages": [AIMessage(content="Actual Worker Findings")],
|
|
"vfs": {},
|
|
"plots": []
|
|
}
|
|
|
|
# Manually replace the runnable with a mock object that has an invoke method
|
|
mock_runnable = MagicMock()
|
|
mock_runnable.invoke = mock_invoke
|
|
monkeypatch.setattr("ea_chatbot.graph.workflow.data_analyst_worker_runnable", mock_runnable)
|
|
|
|
# Execute via the module reference (which is now mocked)
|
|
from ea_chatbot.graph.workflow import data_analyst_worker_runnable
|
|
updates = data_analyst_worker_runnable.invoke(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"
|