52 lines
1.7 KiB
Python
52 lines
1.7 KiB
Python
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
|