forked from cleveragents/cleveragents-core
1375 lines
49 KiB
Python
1375 lines
49 KiB
Python
"""
|
|
Comprehensive unit tests for langgraph graph module.
|
|
"""
|
|
|
|
import asyncio
|
|
import tempfile
|
|
from pathlib import Path
|
|
from unittest.mock import AsyncMock, Mock, patch
|
|
|
|
import pytest
|
|
|
|
from cleveragents.agents.base import Agent
|
|
from cleveragents.langgraph.graph import GraphConfig, LangGraph
|
|
from cleveragents.langgraph.nodes import Edge, NodeConfig, NodeType
|
|
from cleveragents.langgraph.state import GraphState
|
|
from cleveragents.reactive.stream_router import ReactiveStreamRouter, StreamMessage
|
|
from rx.scheduler.eventloop import AsyncIOScheduler
|
|
|
|
|
|
@pytest.fixture
|
|
def linear_graph_structure():
|
|
"""Fixture for a simple linear graph structure (start -> node1 -> end)."""
|
|
nodes = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
"node1": NodeConfig(name="node1", type=NodeType.AGENT),
|
|
"end": NodeConfig(name="end", type=NodeType.END),
|
|
}
|
|
edges = [
|
|
Edge(source="start", target="node1"),
|
|
Edge(source="node1", target="end"),
|
|
]
|
|
return nodes, edges
|
|
|
|
|
|
@pytest.fixture
|
|
def minimal_graph_structure():
|
|
"""Fixture for a minimal graph structure (start -> end)."""
|
|
nodes = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
"end": NodeConfig(name="end", type=NodeType.END),
|
|
}
|
|
edges = [Edge(source="start", target="end")]
|
|
return nodes, edges
|
|
|
|
|
|
@pytest.fixture
|
|
def start_to_node1_structure():
|
|
"""Fixture for a simple graph structure (start -> node1)."""
|
|
nodes = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
"node1": NodeConfig(name="node1", type=NodeType.AGENT),
|
|
}
|
|
edges = [Edge(source="start", target="node1")]
|
|
return nodes, edges
|
|
|
|
|
|
@pytest.fixture
|
|
def parallel_graph_structure():
|
|
"""Fixture for a graph structure with parallel nodes (start -> node1, node2 -> end)."""
|
|
nodes = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
"node1": NodeConfig(name="node1", type=NodeType.AGENT),
|
|
"node2": NodeConfig(name="node2", type=NodeType.AGENT),
|
|
"end": NodeConfig(name="end", type=NodeType.END),
|
|
}
|
|
edges = [
|
|
Edge(source="start", target="node1"),
|
|
Edge(source="start", target="node2"),
|
|
Edge(source="node1", target="end"),
|
|
Edge(source="node2", target="end"),
|
|
]
|
|
return nodes, edges
|
|
|
|
|
|
@pytest.fixture
|
|
def conditional_edge_structure():
|
|
"""Fixture for a graph structure with conditional edge (start -> node1 with condition)."""
|
|
nodes = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
"node1": NodeConfig(name="node1", type=NodeType.AGENT),
|
|
}
|
|
edges = [
|
|
Edge(
|
|
source="start",
|
|
target="node1",
|
|
condition={"type": "always"}
|
|
)
|
|
]
|
|
return nodes, edges
|
|
|
|
|
|
@pytest.fixture
|
|
def cycle_graph_structure():
|
|
"""Fixture for a graph structure with cycles (start -> node1 -> node2 -> node1)."""
|
|
nodes = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
"node1": NodeConfig(name="node1", type=NodeType.AGENT),
|
|
"node2": NodeConfig(name="node2", type=NodeType.AGENT),
|
|
}
|
|
edges = [
|
|
Edge(source="start", target="node1"),
|
|
Edge(source="node1", target="node2"),
|
|
Edge(source="node2", target="node1"),
|
|
]
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges)
|
|
graph = LangGraph(config)
|
|
return graph
|
|
|
|
|
|
@pytest.fixture
|
|
def node_stream_setup():
|
|
"""Fixture for setting up a graph with a single function node for stream subscription tests."""
|
|
nodes = {
|
|
"node1": NodeConfig(name="node1", type=NodeType.FUNCTION, function="test")
|
|
}
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=[])
|
|
agents = {}
|
|
stream_router = ReactiveStreamRouter()
|
|
graph = LangGraph(config, agents, stream_router)
|
|
stream_name = "__test_graph_node_node1__"
|
|
return graph, stream_router, stream_name
|
|
|
|
|
|
class TestGraphConfig:
|
|
"""Test cases for GraphConfig dataclass."""
|
|
|
|
def test_graph_config_defaults(self):
|
|
"""Test GraphConfig with default values."""
|
|
config = GraphConfig(name="test_graph")
|
|
|
|
assert config.name == "test_graph"
|
|
assert config.nodes == {}
|
|
assert config.edges == []
|
|
assert config.entry_point == "start"
|
|
assert config.state_class is None
|
|
assert config.checkpointing is False
|
|
assert config.checkpoint_dir is None
|
|
assert config.enable_time_travel is False
|
|
assert config.parallel_execution is True
|
|
assert config.metadata == {}
|
|
|
|
def test_graph_config_with_values(self):
|
|
"""Test GraphConfig with custom values."""
|
|
nodes = {"node1": NodeConfig(name="node1", type=NodeType.AGENT)}
|
|
edges = [Edge(source="start", target="node1")]
|
|
|
|
config = GraphConfig(
|
|
name="custom_graph",
|
|
nodes=nodes,
|
|
edges=edges,
|
|
entry_point="start",
|
|
checkpointing=True,
|
|
parallel_execution=False,
|
|
metadata={"key": "value"}
|
|
)
|
|
|
|
assert config.name == "custom_graph"
|
|
assert len(config.nodes) == 1
|
|
assert len(config.edges) == 1
|
|
assert config.checkpointing is True
|
|
assert config.parallel_execution is False
|
|
assert config.metadata["key"] == "value"
|
|
|
|
|
|
class TestLangGraphInit:
|
|
"""Test cases for LangGraph initialization."""
|
|
|
|
def test_langgraph_init_minimal(self):
|
|
"""Test LangGraph initialization with minimal config."""
|
|
config = GraphConfig(name="test_graph")
|
|
graph = LangGraph(config)
|
|
|
|
assert graph.name == "test_graph"
|
|
assert graph.config == config
|
|
assert "start" in graph.nodes
|
|
assert "end" in graph.nodes
|
|
assert graph.state_manager is not None
|
|
assert not graph.is_running
|
|
|
|
def test_langgraph_init_with_agents(self):
|
|
"""Test LangGraph initialization with agents."""
|
|
config = GraphConfig(name="test_graph")
|
|
mock_agent = Mock(spec=Agent)
|
|
agents = {"agent1": mock_agent}
|
|
|
|
graph = LangGraph(config, agents=agents)
|
|
|
|
assert "agent1" in graph.agents
|
|
assert graph.agents["agent1"] == mock_agent
|
|
|
|
def test_langgraph_init_with_custom_nodes(self):
|
|
"""Test LangGraph initialization with custom nodes."""
|
|
nodes = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
"agent1": NodeConfig(name="agent1", type=NodeType.AGENT),
|
|
"end": NodeConfig(name="end", type=NodeType.END),
|
|
}
|
|
config = GraphConfig(name="test_graph", nodes=nodes)
|
|
|
|
graph = LangGraph(config)
|
|
|
|
assert "agent1" in graph.nodes
|
|
assert graph.nodes["agent1"].type == NodeType.AGENT
|
|
|
|
def test_langgraph_init_adds_start_end_nodes(self):
|
|
"""Test that LangGraph automatically adds start and end nodes."""
|
|
config = GraphConfig(name="test_graph")
|
|
graph = LangGraph(config)
|
|
|
|
assert "start" in graph.nodes
|
|
assert "end" in graph.nodes
|
|
assert graph.nodes["start"].type == NodeType.START
|
|
assert graph.nodes["end"].type == NodeType.END
|
|
|
|
def test_langgraph_init_with_checkpointing(self):
|
|
"""Test LangGraph initialization with checkpointing."""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
checkpoint_dir = Path(tmpdir)
|
|
config = GraphConfig(
|
|
name="test_graph",
|
|
checkpointing=True,
|
|
checkpoint_dir=checkpoint_dir
|
|
)
|
|
|
|
graph = LangGraph(config)
|
|
|
|
assert graph.state_manager.checkpoint_dir == checkpoint_dir
|
|
|
|
def test_langgraph_init_with_time_travel(self):
|
|
"""Test LangGraph initialization with time travel."""
|
|
config = GraphConfig(name="test_graph", enable_time_travel=True)
|
|
graph = LangGraph(config)
|
|
|
|
assert graph.state_manager.enable_time_travel is True
|
|
|
|
def test_langgraph_init_with_custom_state_class(self):
|
|
"""Test LangGraph initialization with custom state class."""
|
|
class CustomState(GraphState):
|
|
pass
|
|
|
|
config = GraphConfig(name="test_graph", state_class=CustomState)
|
|
graph = LangGraph(config)
|
|
|
|
assert isinstance(graph.state_manager.state, CustomState)
|
|
|
|
def test_langgraph_init_creates_scheduler(self):
|
|
"""Test that LangGraph creates scheduler if not provided."""
|
|
config = GraphConfig(name="test_graph")
|
|
graph = LangGraph(config)
|
|
|
|
assert graph.scheduler is not None
|
|
|
|
def test_langgraph_init_uses_provided_scheduler(self):
|
|
"""Test that LangGraph uses provided scheduler."""
|
|
config = GraphConfig(name="test_graph")
|
|
mock_scheduler = Mock()
|
|
|
|
graph = LangGraph(config, scheduler=mock_scheduler)
|
|
|
|
assert graph.scheduler == mock_scheduler
|
|
|
|
def test_langgraph_init_creates_stream_router(self):
|
|
"""Test that LangGraph creates stream router if not provided."""
|
|
config = GraphConfig(name="test_graph")
|
|
graph = LangGraph(config)
|
|
|
|
assert graph.stream_router is not None
|
|
assert isinstance(graph.stream_router, ReactiveStreamRouter)
|
|
assert hasattr(graph.stream_router, 'streams')
|
|
assert hasattr(graph.stream_router, 'observables')
|
|
assert hasattr(graph.stream_router, 'agents')
|
|
|
|
def test_langgraph_init_uses_provided_stream_router(self):
|
|
"""Test that LangGraph uses provided stream router."""
|
|
config = GraphConfig(name="test_graph")
|
|
mock_router = Mock(spec=ReactiveStreamRouter)
|
|
mock_router.streams = {}
|
|
mock_router.observables = {}
|
|
|
|
with patch.object(LangGraph, '_create_graph_streams'):
|
|
graph = LangGraph(config, stream_router=mock_router)
|
|
assert graph.stream_router == mock_router
|
|
|
|
|
|
class TestLangGraphAnalysis:
|
|
"""Test cases for graph analysis methods."""
|
|
|
|
def test_analyze_graph_builds_adjacency_lists(self, linear_graph_structure):
|
|
"""Test that _analyze_graph builds adjacency lists."""
|
|
nodes, edges = linear_graph_structure
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges)
|
|
|
|
graph = LangGraph(config)
|
|
|
|
assert "node1" in graph.adjacency_list["start"]
|
|
assert "end" in graph.adjacency_list["node1"]
|
|
assert "start" in graph.reverse_adjacency_list["node1"]
|
|
|
|
def test_detect_cycles_no_cycles(self, linear_graph_structure):
|
|
"""Test cycle detection with no cycles."""
|
|
nodes, edges = linear_graph_structure
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges)
|
|
|
|
graph = LangGraph(config)
|
|
|
|
assert graph.has_cycles is False
|
|
|
|
def test_detect_cycles_with_cycles(self, cycle_graph_structure):
|
|
"""Test cycle detection with cycles."""
|
|
graph = cycle_graph_structure
|
|
|
|
assert graph.has_cycles is True
|
|
|
|
def test_find_parallel_groups_disabled(self):
|
|
"""Test parallel group finding when disabled."""
|
|
config = GraphConfig(name="test_graph", parallel_execution=False)
|
|
graph = LangGraph(config)
|
|
|
|
assert graph.parallel_groups == []
|
|
|
|
def test_find_parallel_groups_enabled(self, parallel_graph_structure):
|
|
"""Test parallel group finding when enabled."""
|
|
nodes, edges = parallel_graph_structure
|
|
# Create new NodeConfig objects with parallel=True
|
|
nodes["node1"] = NodeConfig(name="node1", type=NodeType.AGENT, parallel=True)
|
|
nodes["node2"] = NodeConfig(name="node2", type=NodeType.AGENT, parallel=True)
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges, parallel_execution=True)
|
|
|
|
graph = LangGraph(config)
|
|
|
|
# node1 and node2 can execute in parallel - should find at least one group
|
|
assert len(graph.parallel_groups) > 0
|
|
# Verify that node1 and node2 are in a parallel group together
|
|
found_group = any("node1" in group and "node2" in group for group in graph.parallel_groups)
|
|
assert found_group
|
|
|
|
def test_find_parallel_groups_with_multiple_nodes(self):
|
|
"""Test parallel group finding where len(parallel_nodes) > 1."""
|
|
nodes = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
"node1": NodeConfig(name="node1", type=NodeType.AGENT, parallel=True),
|
|
"node2": NodeConfig(name="node2", type=NodeType.AGENT, parallel=True),
|
|
"node3": NodeConfig(name="node3", type=NodeType.AGENT, parallel=True),
|
|
"end": NodeConfig(name="end", type=NodeType.END),
|
|
}
|
|
edges = [
|
|
Edge(source="start", target="node1"),
|
|
Edge(source="start", target="node2"),
|
|
Edge(source="start", target="node3"),
|
|
Edge(source="node1", target="end"),
|
|
Edge(source="node2", target="end"),
|
|
Edge(source="node3", target="end"),
|
|
]
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges, parallel_execution=True)
|
|
|
|
graph = LangGraph(config)
|
|
|
|
# Should find at least one parallel group with multiple nodes
|
|
assert len(graph.parallel_groups) > 0
|
|
has_multi_node_group = any(len(group) > 1 for group in graph.parallel_groups)
|
|
assert has_multi_node_group
|
|
# Verify all three nodes are in the same parallel group
|
|
found_all_nodes = any(
|
|
"node1" in group and "node2" in group and "node3" in group
|
|
for group in graph.parallel_groups
|
|
)
|
|
assert found_all_nodes
|
|
|
|
def test_topological_levels(self):
|
|
"""Test topological level computation."""
|
|
nodes = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
"node1": NodeConfig(name="node1", type=NodeType.AGENT),
|
|
"node2": NodeConfig(name="node2", type=NodeType.AGENT),
|
|
"end": NodeConfig(name="end", type=NodeType.END),
|
|
}
|
|
edges = [
|
|
Edge(source="start", target="node1"),
|
|
Edge(source="node1", target="node2"),
|
|
Edge(source="node2", target="end"),
|
|
]
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges)
|
|
|
|
graph = LangGraph(config)
|
|
levels = graph._topological_levels()
|
|
|
|
assert isinstance(levels, dict)
|
|
# Should have multiple levels for this graph structure
|
|
assert len(levels) >= 2
|
|
# Verify start is at level 0
|
|
assert "start" in levels[0]
|
|
# Verify end is at the last level
|
|
max_level = max(levels.keys())
|
|
assert "end" in levels[max_level]
|
|
# Verify node1 and node2 are at intermediate levels
|
|
assert any("node1" in levels[level] for level in levels)
|
|
assert any("node2" in levels[level] for level in levels)
|
|
|
|
def test_validate_graph_invalid_entry_point(self):
|
|
"""Test graph validation with invalid entry point."""
|
|
config = GraphConfig(name="test_graph", entry_point="nonexistent")
|
|
|
|
with pytest.raises(ValueError, match="Entry point.*not found"):
|
|
LangGraph(config)
|
|
|
|
def test_validate_graph_invalid_edge_source(self):
|
|
"""Test graph validation with invalid edge source."""
|
|
edges = [Edge(source="nonexistent", target="end")]
|
|
config = GraphConfig(name="test_graph", edges=edges)
|
|
|
|
with pytest.raises(ValueError, match="Edge source.*not found"):
|
|
LangGraph(config)
|
|
|
|
def test_validate_graph_invalid_edge_target(self):
|
|
"""Test graph validation with invalid edge target."""
|
|
nodes = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
}
|
|
edges = [Edge(source="start", target="nonexistent")]
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges)
|
|
|
|
with pytest.raises((ValueError, KeyError)):
|
|
LangGraph(config)
|
|
|
|
def test_validate_graph_unreachable_nodes(self):
|
|
"""Test graph validation with unreachable nodes."""
|
|
nodes = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
"unreachable": NodeConfig(name="unreachable", type=NodeType.AGENT),
|
|
"end": NodeConfig(name="end", type=NodeType.END),
|
|
}
|
|
edges = [Edge(source="start", target="end")]
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges)
|
|
|
|
# Should log warning but not fail
|
|
graph = LangGraph(config)
|
|
assert "unreachable" in graph.nodes
|
|
|
|
def test_find_reachable_nodes(self):
|
|
"""Test finding reachable nodes from start."""
|
|
nodes = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
"node1": NodeConfig(name="node1", type=NodeType.AGENT),
|
|
"node2": NodeConfig(name="node2", type=NodeType.AGENT),
|
|
"unreachable": NodeConfig(name="unreachable", type=NodeType.AGENT),
|
|
}
|
|
edges = [
|
|
Edge(source="start", target="node1"),
|
|
Edge(source="node1", target="node2"),
|
|
]
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges)
|
|
|
|
graph = LangGraph(config)
|
|
reachable = graph._find_reachable_nodes("start")
|
|
|
|
assert "start" in reachable
|
|
assert "node1" in reachable
|
|
assert "node2" in reachable
|
|
assert "unreachable" not in reachable
|
|
|
|
def test_find_reachable_nodes_with_cycle(self, cycle_graph_structure):
|
|
"""Test finding reachable nodes when graph has cycles (node in visited branch)."""
|
|
graph = cycle_graph_structure
|
|
reachable = graph._find_reachable_nodes("start")
|
|
|
|
# All nodes should be reachable despite the cycle
|
|
assert "start" in reachable
|
|
assert "node1" in reachable
|
|
assert "node2" in reachable
|
|
# Should handle the cycle without infinite loop (node in visited check)
|
|
|
|
|
|
class TestLangGraphExecution:
|
|
"""Test cases for graph execution."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_simple_graph(self, minimal_graph_structure):
|
|
"""Test executing a simple graph."""
|
|
nodes, edges = minimal_graph_structure
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges)
|
|
|
|
graph = LangGraph(config)
|
|
result = await graph.execute()
|
|
|
|
assert isinstance(result, GraphState)
|
|
assert not graph.is_running
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_with_input_data_messages(self):
|
|
"""Test executing with input data containing messages."""
|
|
config = GraphConfig(name="test_graph")
|
|
graph = LangGraph(config)
|
|
|
|
input_data = {
|
|
"messages": [{"role": "user", "content": "Hello"}]
|
|
}
|
|
result = await graph.execute(input_data)
|
|
|
|
assert result.messages[0]["content"] == "Hello"
|
|
assert result.messages[0]["role"] == "user"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_with_input_data_wrapped(self):
|
|
"""Test executing with input data that gets wrapped as message."""
|
|
config = GraphConfig(name="test_graph")
|
|
graph = LangGraph(config)
|
|
|
|
input_data = {"key": "value"}
|
|
result = await graph.execute(input_data)
|
|
|
|
# Input should be wrapped as a message
|
|
assert len(result.messages) == 1
|
|
assert result.messages[0]["content"] == {"key": "value"}
|
|
assert result.messages[0]["role"] == "user"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_with_metadata(self):
|
|
"""Test executing with metadata in input."""
|
|
config = GraphConfig(name="test_graph")
|
|
graph = LangGraph(config)
|
|
|
|
input_data = {
|
|
"metadata": {"test_key": "test_value"}
|
|
}
|
|
result = await graph.execute(input_data)
|
|
|
|
assert "test_key" in result.metadata
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_already_running(self):
|
|
"""Test that execute raises error if already running."""
|
|
config = GraphConfig(name="test_graph")
|
|
graph = LangGraph(config)
|
|
graph.is_running = True
|
|
|
|
with pytest.raises(RuntimeError, match="already running"):
|
|
await graph.execute()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_clears_history(self):
|
|
"""Test that execute clears execution history."""
|
|
config = GraphConfig(name="test_graph")
|
|
graph = LangGraph(config)
|
|
graph.execution_history = ["old_node"]
|
|
|
|
await graph.execute()
|
|
|
|
# History should be cleared at start
|
|
assert "old_node" not in graph.execution_history or len(graph.execution_history) > 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_from_node_sequential(self, minimal_graph_structure):
|
|
"""Test executing from a node in sequential mode."""
|
|
nodes, edges = minimal_graph_structure
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges, parallel_execution=False)
|
|
|
|
graph = LangGraph(config)
|
|
await graph._execute_from_node("start")
|
|
|
|
# Should have executed at least the start node
|
|
assert len(graph.execution_history) > 0
|
|
assert "start" in graph.execution_history
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_from_node_parallel(self, parallel_graph_structure):
|
|
"""Test executing from a node in parallel mode."""
|
|
nodes, edges = parallel_graph_structure
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges, parallel_execution=True)
|
|
|
|
graph = LangGraph(config)
|
|
await graph._execute_from_node("start")
|
|
|
|
# Should have executed at least the start node
|
|
assert len(graph.execution_history) > 0
|
|
assert "start" in graph.execution_history
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_can_execute_node_no_predecessors(self):
|
|
"""Test _can_execute_node with no predecessors."""
|
|
config = GraphConfig(name="test_graph")
|
|
graph = LangGraph(config)
|
|
|
|
# Start node has no predecessors
|
|
result = graph._can_execute_node("start", set())
|
|
|
|
assert result is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_can_execute_node_with_executed_predecessors(self, start_to_node1_structure):
|
|
"""Test _can_execute_node with executed predecessors."""
|
|
nodes, edges = start_to_node1_structure
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges)
|
|
|
|
graph = LangGraph(config)
|
|
executed = {"start"}
|
|
|
|
result = graph._can_execute_node("node1", executed)
|
|
|
|
assert result is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_can_execute_node_with_unexecuted_predecessors(self, start_to_node1_structure):
|
|
"""Test _can_execute_node with unexecuted predecessors."""
|
|
nodes, edges = start_to_node1_structure
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges)
|
|
|
|
graph = LangGraph(config)
|
|
executed = set() # Start not executed
|
|
|
|
result = graph._can_execute_node("node1", executed)
|
|
|
|
assert result is False
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_node(self):
|
|
"""Test executing a single node."""
|
|
config = GraphConfig(name="test_graph")
|
|
graph = LangGraph(config)
|
|
|
|
await graph._execute_node("start")
|
|
|
|
# Should have sent message to stream
|
|
assert "start" in graph.execution_history
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_nodes_parallel(self):
|
|
"""Test executing multiple nodes in parallel."""
|
|
nodes = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
"node1": NodeConfig(name="node1", type=NodeType.AGENT),
|
|
"node2": NodeConfig(name="node2", type=NodeType.AGENT),
|
|
}
|
|
config = GraphConfig(name="test_graph", nodes=nodes)
|
|
|
|
graph = LangGraph(config)
|
|
await graph._execute_nodes_parallel(["start", "node1", "node2"])
|
|
|
|
# Should have executed all three nodes
|
|
assert len(graph.execution_history) == 3
|
|
assert "start" in graph.execution_history
|
|
assert "node1" in graph.execution_history
|
|
assert "node2" in graph.execution_history
|
|
|
|
def test_get_next_nodes_no_edges(self):
|
|
"""Test getting next nodes when there are no outgoing edges."""
|
|
config = GraphConfig(name="test_graph")
|
|
graph = LangGraph(config)
|
|
|
|
next_nodes = graph._get_next_nodes("end")
|
|
|
|
assert next_nodes == []
|
|
|
|
def test_get_next_nodes_with_edges(self, start_to_node1_structure):
|
|
"""Test getting next nodes with outgoing edges."""
|
|
nodes, edges = start_to_node1_structure
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges)
|
|
|
|
graph = LangGraph(config)
|
|
next_nodes = graph._get_next_nodes("start")
|
|
|
|
assert "node1" in next_nodes
|
|
|
|
def test_get_next_nodes_with_condition(self, conditional_edge_structure):
|
|
"""Test getting next nodes with conditional edges."""
|
|
nodes, edges = conditional_edge_structure
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges)
|
|
|
|
graph = LangGraph(config)
|
|
next_nodes = graph._get_next_nodes("start")
|
|
|
|
assert "node1" in next_nodes
|
|
|
|
|
|
class TestLangGraphMutation:
|
|
"""Test cases for graph mutation methods."""
|
|
|
|
def test_add_node(self):
|
|
"""Test adding a node to the graph."""
|
|
config = GraphConfig(name="test_graph")
|
|
graph = LangGraph(config)
|
|
|
|
new_node = NodeConfig(name="new_node", type=NodeType.AGENT)
|
|
graph.add_node(new_node)
|
|
|
|
assert "new_node" in graph.nodes
|
|
assert "new_node" in graph.config.nodes
|
|
|
|
def test_add_node_duplicate(self):
|
|
"""Test adding a duplicate node raises error."""
|
|
config = GraphConfig(name="test_graph")
|
|
graph = LangGraph(config)
|
|
|
|
new_node = NodeConfig(name="start", type=NodeType.START)
|
|
|
|
with pytest.raises(ValueError, match="already exists"):
|
|
graph.add_node(new_node)
|
|
|
|
def test_add_edge(self):
|
|
"""Test adding an edge to the graph."""
|
|
nodes = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
"node1": NodeConfig(name="node1", type=NodeType.AGENT),
|
|
}
|
|
config = GraphConfig(name="test_graph", nodes=nodes)
|
|
graph = LangGraph(config)
|
|
|
|
new_edge = Edge(source="start", target="node1")
|
|
graph.add_edge(new_edge)
|
|
|
|
assert new_edge in graph.config.edges
|
|
|
|
def test_add_edge_invalid_source(self):
|
|
"""Test adding edge with invalid source raises error."""
|
|
config = GraphConfig(name="test_graph")
|
|
graph = LangGraph(config)
|
|
|
|
new_edge = Edge(source="nonexistent", target="end")
|
|
|
|
with pytest.raises(ValueError, match="Source node.*not found"):
|
|
graph.add_edge(new_edge)
|
|
|
|
def test_add_edge_invalid_target(self):
|
|
"""Test adding edge with invalid target raises error."""
|
|
config = GraphConfig(name="test_graph")
|
|
graph = LangGraph(config)
|
|
|
|
new_edge = Edge(source="start", target="nonexistent")
|
|
|
|
with pytest.raises(ValueError, match="Target node.*not found"):
|
|
graph.add_edge(new_edge)
|
|
|
|
|
|
class TestLangGraphAccessors:
|
|
"""Test cases for graph accessor methods."""
|
|
|
|
def test_get_state(self):
|
|
"""Test getting current graph state."""
|
|
config = GraphConfig(name="test_graph")
|
|
graph = LangGraph(config)
|
|
|
|
state = graph.get_state()
|
|
|
|
assert isinstance(state, GraphState)
|
|
|
|
def test_get_execution_history(self):
|
|
"""Test getting execution history."""
|
|
config = GraphConfig(name="test_graph")
|
|
graph = LangGraph(config)
|
|
graph.execution_history = ["node1", "node2"]
|
|
|
|
history = graph.get_execution_history()
|
|
|
|
assert history == ["node1", "node2"]
|
|
# Should be a copy, not the original list
|
|
history.append("node3")
|
|
assert len(graph.execution_history) == 2
|
|
|
|
|
|
class TestLangGraphVisualization:
|
|
"""Test cases for graph visualization."""
|
|
|
|
def test_visualize_mermaid_format(self):
|
|
"""Test visualizing graph in mermaid format."""
|
|
nodes = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
"agent": NodeConfig(name="agent", type=NodeType.AGENT),
|
|
"end": NodeConfig(name="end", type=NodeType.END),
|
|
}
|
|
edges = [
|
|
Edge(source="start", target="agent"),
|
|
Edge(source="agent", target="end"),
|
|
]
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges)
|
|
|
|
graph = LangGraph(config)
|
|
result = graph.visualize()
|
|
|
|
assert "graph TD" in result
|
|
assert "start" in result
|
|
assert "agent" in result
|
|
assert "end" in result
|
|
assert "-->" in result
|
|
|
|
def test_visualize_with_conditional_edges(self, conditional_edge_structure):
|
|
"""Test visualizing graph with conditional edges."""
|
|
nodes, edges = conditional_edge_structure
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges)
|
|
|
|
graph = LangGraph(config)
|
|
result = graph.visualize()
|
|
|
|
assert "always" in result or "condition" in result
|
|
|
|
def test_visualize_unsupported_format(self):
|
|
"""Test visualizing with unsupported format."""
|
|
config = GraphConfig(name="test_graph")
|
|
graph = LangGraph(config)
|
|
|
|
result = graph.visualize(output_format="graphviz")
|
|
|
|
assert "not supported" in result
|
|
|
|
def test_get_node_shape_all_types(self):
|
|
"""Test getting node shapes for all node types."""
|
|
config = GraphConfig(name="test_graph")
|
|
graph = LangGraph(config)
|
|
|
|
shapes = {
|
|
NodeType.START: "((Start))",
|
|
NodeType.END: "((End))",
|
|
NodeType.AGENT: "[Agent]",
|
|
NodeType.FUNCTION: "[Function]",
|
|
NodeType.TOOL: "[/Tool/]",
|
|
NodeType.CONDITIONAL: "{Conditional}",
|
|
NodeType.SUBGRAPH: "[[Subgraph]]",
|
|
}
|
|
|
|
for node_type, expected_shape in shapes.items():
|
|
shape = graph._get_node_shape(node_type)
|
|
assert shape == expected_shape
|
|
|
|
|
|
class TestLangGraphIntegration:
|
|
"""Integration test cases for LangGraph."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_full_graph_execution(self):
|
|
"""Test full graph execution from start to end."""
|
|
nodes = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
"function": NodeConfig(
|
|
name="function",
|
|
type=NodeType.FUNCTION,
|
|
function="summarize"
|
|
),
|
|
"end": NodeConfig(name="end", type=NodeType.END),
|
|
}
|
|
edges = [
|
|
Edge(source="start", target="function"),
|
|
Edge(source="function", target="end"),
|
|
]
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges)
|
|
|
|
graph = LangGraph(config)
|
|
result = await graph.execute({"messages": [{"role": "user", "content": "test"}]})
|
|
|
|
assert isinstance(result, GraphState)
|
|
assert not graph.is_running
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_graph_with_conditional_routing(self):
|
|
"""Test graph execution with conditional routing."""
|
|
nodes = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
"node1": NodeConfig(name="node1", type=NodeType.AGENT),
|
|
"node2": NodeConfig(name="node2", type=NodeType.AGENT),
|
|
"end": NodeConfig(name="end", type=NodeType.END),
|
|
}
|
|
edges = [
|
|
Edge(source="start", target="node1", condition={"type": "always"}),
|
|
Edge(source="node1", target="end"),
|
|
]
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges)
|
|
|
|
graph = LangGraph(config)
|
|
result = await graph.execute()
|
|
|
|
assert isinstance(result, GraphState)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_graph_with_agent_node(self):
|
|
"""Test graph execution with agent node."""
|
|
mock_agent = AsyncMock(spec=Agent)
|
|
mock_agent.process_message = AsyncMock(return_value={
|
|
"role": "assistant",
|
|
"content": "Response"
|
|
})
|
|
|
|
nodes = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
"agent": NodeConfig(
|
|
name="agent",
|
|
type=NodeType.AGENT,
|
|
agent="test_agent"
|
|
),
|
|
"end": NodeConfig(name="end", type=NodeType.END),
|
|
}
|
|
edges = [
|
|
Edge(source="start", target="agent"),
|
|
Edge(source="agent", target="end"),
|
|
]
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges)
|
|
|
|
agents = {"test_agent": mock_agent}
|
|
graph = LangGraph(config, agents=agents)
|
|
result = await graph.execute({"messages": [{"role": "user", "content": "test"}]})
|
|
|
|
assert isinstance(result, GraphState)
|
|
|
|
def test_graph_state_persistence(self):
|
|
"""Test that graph state is persisted across operations."""
|
|
config = GraphConfig(name="test_graph")
|
|
graph = LangGraph(config)
|
|
|
|
# Update state
|
|
graph.state_manager.update_state({"current_node": "test"})
|
|
|
|
# Get state
|
|
state = graph.get_state()
|
|
|
|
assert state.current_node == "test"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_graph_execution_with_time_travel(self):
|
|
"""Test graph execution with time travel enabled."""
|
|
config = GraphConfig(name="test_graph", enable_time_travel=True)
|
|
graph = LangGraph(config)
|
|
|
|
await graph.execute()
|
|
|
|
# Should have history
|
|
assert graph.state_manager.enable_time_travel is True
|
|
|
|
|
|
class TestLangGraphEdgeCases:
|
|
"""Test cases for edge cases and corner scenarios."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_from_node_no_progress(self):
|
|
"""Test _execute_from_node when no progress can be made."""
|
|
nodes = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
"node1": NodeConfig(name="node1", type=NodeType.AGENT),
|
|
"node2": NodeConfig(name="node2", type=NodeType.AGENT),
|
|
}
|
|
# Create edges with conditions that won't be met
|
|
edges = [
|
|
Edge(source="start", target="node1", condition={"type": "never"}),
|
|
Edge(source="node1", target="node2"),
|
|
]
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges)
|
|
|
|
graph = LangGraph(config)
|
|
await graph._execute_from_node("start")
|
|
|
|
# Should still complete without hanging
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_can_execute_node_with_conditional_edge(self):
|
|
"""Test _can_execute_node with conditional edge evaluation."""
|
|
nodes = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
"node1": NodeConfig(name="node1", type=NodeType.AGENT),
|
|
}
|
|
edges = [
|
|
Edge(source="start", target="node1", condition={"type": "never"})
|
|
]
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges)
|
|
|
|
graph = LangGraph(config)
|
|
executed = {"start"}
|
|
|
|
result = graph._can_execute_node("node1", executed)
|
|
|
|
# Should be False because condition won't be satisfied
|
|
assert result is False
|
|
|
|
def test_adjacency_list_empty_graph(self):
|
|
"""Test adjacency lists with no edges."""
|
|
config = GraphConfig(name="test_graph")
|
|
graph = LangGraph(config)
|
|
|
|
# Should have empty adjacency for most nodes
|
|
assert isinstance(graph.adjacency_list, dict)
|
|
assert isinstance(graph.reverse_adjacency_list, dict)
|
|
|
|
def test_topological_levels_single_node(self, minimal_graph_structure):
|
|
"""Test topological levels with single path."""
|
|
nodes, edges = minimal_graph_structure
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges)
|
|
|
|
graph = LangGraph(config)
|
|
levels = graph._topological_levels()
|
|
|
|
# Should have at least 2 levels
|
|
assert len(levels) >= 2
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_with_empty_input(self):
|
|
"""Test execute with None input."""
|
|
config = GraphConfig(name="test_graph")
|
|
graph = LangGraph(config)
|
|
|
|
result = await graph.execute(None)
|
|
|
|
assert isinstance(result, GraphState)
|
|
|
|
def test_parallel_groups_with_single_nodes(self):
|
|
"""Test parallel group finding with nodes that can't parallelize."""
|
|
nodes = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
"node1": NodeConfig(name="node1", type=NodeType.AGENT, parallel=False),
|
|
"end": NodeConfig(name="end", type=NodeType.END),
|
|
}
|
|
edges = [
|
|
Edge(source="start", target="node1"),
|
|
Edge(source="node1", target="end"),
|
|
]
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges, parallel_execution=True)
|
|
|
|
graph = LangGraph(config)
|
|
|
|
# Should not find many parallel groups
|
|
assert isinstance(graph.parallel_groups, list)
|
|
|
|
def test_visualize_edge_without_condition(self, minimal_graph_structure):
|
|
"""Test visualizing edges without conditions."""
|
|
nodes, edges = minimal_graph_structure
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges)
|
|
|
|
graph = LangGraph(config)
|
|
result = graph.visualize()
|
|
|
|
assert "start --> end" in result
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_node_executor_function(self):
|
|
"""Test node executor registration and execution."""
|
|
nodes = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
}
|
|
config = GraphConfig(name="test_graph", nodes=nodes)
|
|
|
|
graph = LangGraph(config)
|
|
|
|
# Verify executor is registered
|
|
assert hasattr(graph.stream_router, '_builtin_execute_node_start')
|
|
|
|
def test_get_node_shape_unknown_type(self):
|
|
"""Test get_node_shape with unknown node type (edge case)."""
|
|
config = GraphConfig(name="test_graph")
|
|
graph = LangGraph(config)
|
|
|
|
# Test that method returns default shape for a NodeType value not in shapes dict
|
|
# Since all NodeType enum values are handled, we test with a mock
|
|
# In practice, this shouldn't happen as method signature requires NodeType
|
|
mock_node_type = Mock(spec=NodeType)
|
|
mock_node_type.value = "unknown"
|
|
|
|
shape = graph._get_node_shape(mock_node_type)
|
|
assert shape == "[Node]"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_multiple_times(self):
|
|
"""Test executing graph multiple times."""
|
|
config = GraphConfig(name="test_graph")
|
|
graph = LangGraph(config)
|
|
|
|
result1 = await graph.execute()
|
|
result2 = await graph.execute()
|
|
|
|
assert isinstance(result1, GraphState)
|
|
assert isinstance(result2, GraphState)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_with_complex_metadata(self):
|
|
"""Test execute with complex metadata structure."""
|
|
config = GraphConfig(name="test_graph")
|
|
graph = LangGraph(config)
|
|
|
|
input_data = {
|
|
"metadata": {
|
|
"nested": {
|
|
"key": "value"
|
|
},
|
|
"list": [1, 2, 3]
|
|
}
|
|
}
|
|
result = await graph.execute(input_data)
|
|
|
|
assert "nested" in result.metadata or "list" in result.metadata
|
|
|
|
def test_find_reachable_nodes_complex_graph(self):
|
|
"""Test finding reachable nodes in complex graph."""
|
|
nodes = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
"node1": NodeConfig(name="node1", type=NodeType.AGENT),
|
|
"node2": NodeConfig(name="node2", type=NodeType.AGENT),
|
|
"node3": NodeConfig(name="node3", type=NodeType.AGENT),
|
|
"node4": NodeConfig(name="node4", type=NodeType.AGENT),
|
|
"end": NodeConfig(name="end", type=NodeType.END),
|
|
}
|
|
edges = [
|
|
Edge(source="start", target="node1"),
|
|
Edge(source="node1", target="node2"),
|
|
Edge(source="node1", target="node3"),
|
|
Edge(source="node2", target="end"),
|
|
Edge(source="node3", target="end"),
|
|
# node4 is unreachable
|
|
]
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges)
|
|
|
|
graph = LangGraph(config)
|
|
reachable = graph._find_reachable_nodes("start")
|
|
|
|
assert "start" in reachable
|
|
assert "node1" in reachable
|
|
assert "node2" in reachable
|
|
assert "node3" in reachable
|
|
assert "end" in reachable
|
|
assert "node4" not in reachable
|
|
|
|
def test_detect_cycles_self_loop(self):
|
|
"""Test cycle detection with self-loop."""
|
|
nodes = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
"node1": NodeConfig(name="node1", type=NodeType.AGENT),
|
|
}
|
|
edges = [
|
|
Edge(source="start", target="node1"),
|
|
Edge(source="node1", target="node1"), # Self-loop
|
|
]
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges)
|
|
|
|
graph = LangGraph(config)
|
|
|
|
assert graph.has_cycles is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_graph_with_metadata_in_config(self):
|
|
"""Test graph with metadata in config."""
|
|
config = GraphConfig(
|
|
name="test_graph",
|
|
metadata={"version": "1.0", "author": "test"}
|
|
)
|
|
graph = LangGraph(config)
|
|
|
|
assert graph.config.metadata["version"] == "1.0"
|
|
assert graph.config.metadata["author"] == "test"
|
|
|
|
def test_add_node_reanalyzes_graph(self):
|
|
"""Test that adding node triggers graph reanalysis."""
|
|
config = GraphConfig(name="test_graph")
|
|
graph = LangGraph(config)
|
|
|
|
initial_node_count = len(graph.nodes)
|
|
|
|
new_node = NodeConfig(name="new_node", type=NodeType.AGENT)
|
|
graph.add_node(new_node)
|
|
|
|
assert len(graph.nodes) == initial_node_count + 1
|
|
# Graph should be reanalyzed (adjacency lists updated)
|
|
assert isinstance(graph.adjacency_list, dict)
|
|
|
|
def test_add_edge_reanalyzes_graph(self):
|
|
"""Test that adding edge triggers graph reanalysis."""
|
|
nodes = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
"node1": NodeConfig(name="node1", type=NodeType.AGENT),
|
|
"end": NodeConfig(name="end", type=NodeType.END),
|
|
}
|
|
config = GraphConfig(name="test_graph", nodes=nodes)
|
|
graph = LangGraph(config)
|
|
|
|
initial_edge_count = len(graph.config.edges)
|
|
|
|
new_edge = Edge(source="start", target="node1")
|
|
graph.add_edge(new_edge)
|
|
|
|
assert len(graph.config.edges) == initial_edge_count + 1
|
|
assert "node1" in graph.adjacency_list["start"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_sets_is_running(self):
|
|
"""Test that execute properly sets is_running flag."""
|
|
config = GraphConfig(name="test_graph")
|
|
graph = LangGraph(config)
|
|
|
|
assert not graph.is_running
|
|
|
|
# Execute should set and unset is_running
|
|
await graph.execute()
|
|
|
|
assert not graph.is_running
|
|
|
|
def test_visualize_all_node_types(self):
|
|
"""Test visualizing graph with all node types."""
|
|
nodes = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
"agent": NodeConfig(name="agent", type=NodeType.AGENT),
|
|
"function": NodeConfig(name="function", type=NodeType.FUNCTION),
|
|
"tool": NodeConfig(name="tool", type=NodeType.TOOL),
|
|
"conditional": NodeConfig(name="conditional", type=NodeType.CONDITIONAL),
|
|
"subgraph": NodeConfig(name="subgraph", type=NodeType.SUBGRAPH),
|
|
"end": NodeConfig(name="end", type=NodeType.END),
|
|
}
|
|
config = GraphConfig(name="test_graph", nodes=nodes)
|
|
|
|
graph = LangGraph(config)
|
|
result = graph.visualize()
|
|
|
|
# All nodes should be in visualization
|
|
for node_name in nodes:
|
|
assert node_name in result
|
|
|
|
def test_topological_levels_disconnected_nodes(self):
|
|
"""Test topological levels with disconnected components."""
|
|
nodes = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
"node1": NodeConfig(name="node1", type=NodeType.AGENT),
|
|
"isolated": NodeConfig(name="isolated", type=NodeType.AGENT),
|
|
}
|
|
edges = [
|
|
Edge(source="start", target="node1"),
|
|
]
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges)
|
|
|
|
graph = LangGraph(config)
|
|
levels = graph._topological_levels()
|
|
|
|
# Isolated node should be at level 0 since it has no incoming edges
|
|
assert "isolated" in levels[0]
|
|
|
|
|
|
class TestLangGraphSchedulerCreation:
|
|
"""Test cases for scheduler creation scenarios."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_scheduler_with_running_event_loop(self):
|
|
"""Test scheduler creation when event loop is already running."""
|
|
config = GraphConfig(name="test_graph")
|
|
|
|
# In an async context, event loop should be running
|
|
graph = LangGraph(config)
|
|
|
|
assert graph.scheduler is not None
|
|
assert isinstance(graph.scheduler, AsyncIOScheduler)
|
|
|
|
def test_scheduler_no_event_loop(self):
|
|
"""Test scheduler creation when no event loop exists."""
|
|
config = GraphConfig(name="test_graph")
|
|
|
|
# Create graph in sync context
|
|
graph = LangGraph(config)
|
|
|
|
assert graph.scheduler is not None
|
|
assert isinstance(graph.scheduler, AsyncIOScheduler)
|
|
|
|
|
|
class TestNodeStreamSubscriptionsCallbacks:
|
|
"""Test cases for _setup_node_stream_subscriptions callback functions."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_setup_node_stream_subscriptions_on_next(self):
|
|
"""Test _setup_node_stream_subscriptions on_next callback."""
|
|
nodes = {
|
|
"node1": NodeConfig(name="node1", type=NodeType.FUNCTION, function="test")
|
|
}
|
|
edges = [Edge(source="start", target="node1"), Edge(source="node1", target="end")]
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges)
|
|
|
|
agents = {}
|
|
stream_router = ReactiveStreamRouter()
|
|
|
|
# The on_next callback is set up automatically when nodes are reachable
|
|
# Create graph to set up node streams
|
|
LangGraph(config, agents, stream_router)
|
|
|
|
# Verify node stream exists
|
|
stream_name = "__test_graph_node_node1__"
|
|
assert stream_name in stream_router.observables
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_setup_node_stream_subscriptions_on_error(self, node_stream_setup):
|
|
"""Test _setup_node_stream_subscriptions on_error callback logs errors."""
|
|
graph, stream_router, stream_name = node_stream_setup
|
|
|
|
# Get the node stream
|
|
if stream_name in stream_router.streams:
|
|
stream = stream_router.streams[stream_name]
|
|
|
|
# Trigger an error (this should be logged, not raised)
|
|
with patch.object(graph.logger, 'error') as mock_log:
|
|
stream.on_error(ValueError("Test error"))
|
|
|
|
# Give time for async processing
|
|
await asyncio.sleep(0.05)
|
|
|
|
# Verify error was logged
|
|
# Note: may not be called if error is caught by RxPy
|
|
assert isinstance(mock_log.call_count, int)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_setup_node_stream_subscriptions_on_completed(self, node_stream_setup):
|
|
"""Test _setup_node_stream_subscriptions on_completed callback."""
|
|
_, stream_router, stream_name = node_stream_setup
|
|
|
|
# Get the node stream
|
|
if stream_name in stream_router.streams:
|
|
stream = stream_router.streams[stream_name]
|
|
|
|
# Trigger completion (should not raise error)
|
|
stream.on_completed()
|
|
|
|
await asyncio.sleep(0.05)
|
|
|
|
# If we get here without error, the callback works
|
|
|
|
|
|
class TestRegisterNodeExecutorCallbacks:
|
|
"""Test cases for _register_node_executor callback functions."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_node_executor_async_executor(self):
|
|
"""Test _register_node_executor async_executor callback."""
|
|
|
|
mock_agent = Mock(spec=Agent)
|
|
mock_agent.name = "test_agent"
|
|
|
|
async def mock_process(content):
|
|
return {"response": "test"}
|
|
|
|
mock_agent.process_message = mock_process
|
|
|
|
nodes = {
|
|
"node1": NodeConfig(name="node1", type=NodeType.AGENT, agent="test_agent")
|
|
}
|
|
edges = [Edge(source="start", target="node1")]
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges)
|
|
|
|
agents = {"test_agent": mock_agent}
|
|
stream_router = ReactiveStreamRouter()
|
|
graph = LangGraph(config, agents, stream_router)
|
|
|
|
# Execute the node to trigger async_executor
|
|
result = await graph.execute({"messages": [{"role": "user", "content": "test"}]})
|
|
|
|
# Verify execution happened
|
|
assert isinstance(result, GraphState)
|
|
assert "node1" in graph.execution_history
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_node_executor_sync_executor(self):
|
|
"""Test _register_node_executor sync_executor callback."""
|
|
nodes = {
|
|
"node1": NodeConfig(name="node1", type=NodeType.FUNCTION, function="summarize")
|
|
}
|
|
edges = [Edge(source="start", target="node1")]
|
|
config = GraphConfig(name="test_graph", nodes=nodes, edges=edges)
|
|
|
|
agents = {}
|
|
stream_router = ReactiveStreamRouter()
|
|
_ = LangGraph(config, agents, stream_router)
|
|
|
|
# Check that the sync executor was registered
|
|
executor_name = "_builtin_execute_node_node1"
|
|
assert hasattr(stream_router, executor_name)
|
|
|
|
# Get the executor function
|
|
executor = getattr(stream_router, executor_name)
|
|
|
|
# Execute it synchronously
|
|
test_message = StreamMessage(content={"test": "data"})
|
|
result = executor(test_message)
|
|
|
|
# Verify it returns a StreamMessage
|
|
assert isinstance(result, StreamMessage)
|
|
assert "node" in result.metadata
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"])
|
|
|