import pytest
from ea_chatbot.graph.workers.data_analyst.mapping import merge_worker_output as merge_analyst
from ea_chatbot.graph.workers.researcher.mapping import merge_researcher_output as merge_researcher
from ea_chatbot.graph.workers.data_analyst.state import WorkerState as AnalystState
from ea_chatbot.graph.workers.researcher.state import WorkerState as ResearcherState

def test_analyst_merge_updates_summary():
    """Verify that analyst merge updates the global summary."""
    worker_state = AnalystState(
        result="Actual Worker Result",
        messages=[],
        task="test",
        iterations=1,
        vfs_state={},
        plots=[]
    )
    
    updates = merge_analyst(worker_state)
    assert "summary" in updates
    assert updates["summary"] == "Actual Worker Result"

def test_researcher_merge_updates_summary():
    """Verify that researcher merge updates the global summary."""
    worker_state = ResearcherState(
        result="Actual Research Result",
        messages=[],
        task="test",
        iterations=1,
        queries=[],
        raw_results=[]
    )
    
    updates = merge_researcher(worker_state)
    assert "summary" in updates
    assert updates["summary"] == "Actual Research Result"

def test_merge_handles_none_result():
    """Verify that merge functions handle None results gracefully."""
    worker_state = AnalystState(
        result=None,
        messages=[],
        task="test",
        iterations=1,
        vfs_state={},
        plots=[]
    )
    
    updates = merge_analyst(worker_state)
    assert updates["summary"] is not None
    assert isinstance(updates["messages"][0].content, str)
    assert len(updates["messages"][0].content) > 0