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"