313 lines
11 KiB
Python
313 lines
11 KiB
Python
import asyncio
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
from behave import given, then, when
|
|
from rx.scheduler.eventloop import AsyncIOScheduler
|
|
|
|
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 StreamMessage
|
|
|
|
# Existing steps cover constructor scheduling, execution dispatch, start validation,
|
|
# cycle detection, and node executor state updates.
|
|
|
|
|
|
@given("a langgraph config with layered nodes")
|
|
def step_layered_config(context):
|
|
nodes = {
|
|
"a": NodeConfig(name="a", type=NodeType.FUNCTION),
|
|
"b": NodeConfig(name="b", type=NodeType.FUNCTION),
|
|
"c": NodeConfig(name="c", type=NodeType.FUNCTION),
|
|
}
|
|
edges = [Edge(source="a", target="b"), Edge(source="b", target="c")]
|
|
context.graph = _fresh_graph(
|
|
GraphConfig(name="graph-levels", nodes=nodes, edges=edges)
|
|
)
|
|
|
|
|
|
@when("I compute the topological levels")
|
|
def step_compute_levels(context):
|
|
context.levels = context.graph._topological_levels() # pylint: disable=protected-access
|
|
|
|
|
|
@then("the levels should list each layer")
|
|
def step_assert_levels(context):
|
|
assert context.levels[0] == {"a"}
|
|
assert context.levels[1] == {"b"}
|
|
assert context.levels[2] == {"c"}
|
|
|
|
|
|
@given("a langgraph instance with recorded history")
|
|
def step_history_instance(context):
|
|
config = GraphConfig(name="graph-history")
|
|
context.graph = _fresh_graph(config)
|
|
context.graph.execution_history = ["x", "y"]
|
|
|
|
|
|
@when("I request execution history")
|
|
def step_get_history(context):
|
|
context.history_result = context.graph.get_execution_history()
|
|
|
|
|
|
@then("it should return a copy of the recorded history")
|
|
def step_assert_history_copy(context):
|
|
assert context.history_result == ["x", "y"]
|
|
assert context.history_result is not context.graph.execution_history
|
|
|
|
|
|
@given("a langgraph config with explicit start and end nodes")
|
|
def step_explicit_guards(context):
|
|
start = NodeConfig(name="start", type=NodeType.START, metadata={"guard": True})
|
|
end = NodeConfig(name="end", type=NodeType.END, metadata={"guard": True})
|
|
nodes = {
|
|
"start": start,
|
|
"end": end,
|
|
"worker": NodeConfig(name="worker", type=NodeType.FUNCTION),
|
|
}
|
|
context.graph = _fresh_graph(GraphConfig(name="graph-guards", nodes=nodes))
|
|
|
|
|
|
@when("I construct the LangGraph with that config")
|
|
def step_construct_with_guards(_context):
|
|
# Construction already performed in the given step
|
|
return
|
|
|
|
|
|
@then("initialization should reuse the provided start and end nodes")
|
|
def step_assert_guard_reuse(context):
|
|
assert context.graph.nodes["start"].config is context.graph.config.nodes["start"]
|
|
assert context.graph.nodes["end"].config is context.graph.config.nodes["end"]
|
|
|
|
|
|
@given("a langgraph instance prepared to observe node streams")
|
|
def step_prepare_node_stream(context):
|
|
nodes = {"worker": NodeConfig(name="worker", type=NodeType.FUNCTION)}
|
|
context.graph = _fresh_graph(GraphConfig(name="graph-observe", nodes=nodes))
|
|
context.worker_stream = context.graph.stream_router.streams[
|
|
f"__{context.graph.name}_node_worker__"
|
|
]
|
|
context.graph.logger = MagicMock()
|
|
|
|
|
|
@when("the node stream emits next error and completion events")
|
|
def step_emit_stream_events(context):
|
|
context.worker_stream.on_next("payload")
|
|
context.worker_stream.on_error(ValueError("boom"))
|
|
context.worker_stream.on_completed()
|
|
|
|
|
|
@then("the subscription handlers should be exercised")
|
|
def step_assert_stream_handlers(context):
|
|
context.graph.logger.error.assert_called_once()
|
|
|
|
|
|
@given("a langgraph config with converging branches")
|
|
def step_converging_config(context):
|
|
nodes = {
|
|
"a": NodeConfig(name="a", type=NodeType.FUNCTION),
|
|
"b": NodeConfig(name="b", type=NodeType.FUNCTION),
|
|
"c": NodeConfig(name="c", type=NodeType.FUNCTION),
|
|
}
|
|
edges = [
|
|
Edge(source="start", target="a"),
|
|
Edge(source="start", target="b"),
|
|
Edge(source="a", target="c"),
|
|
Edge(source="b", target="c"),
|
|
Edge(source="c", target="end"),
|
|
]
|
|
context.graph = _fresh_graph(
|
|
GraphConfig(name="graph-parallel", nodes=nodes, edges=edges)
|
|
)
|
|
|
|
|
|
@when("I construct the LangGraph for parallel groups")
|
|
def step_construct_parallel(_context):
|
|
return
|
|
|
|
|
|
@then("parallel groups should reflect staged execution order")
|
|
def step_assert_parallel_groups(context):
|
|
assert context.graph.parallel_groups[0] == {"start"}
|
|
assert context.graph.parallel_groups[1] == {"a", "b"}
|
|
assert context.graph.parallel_groups[2] == {"c"}
|
|
assert context.graph.parallel_groups[3] == {"end"}
|
|
|
|
|
|
@given("a langgraph instance missing required guard nodes")
|
|
def step_missing_guard_nodes(context):
|
|
config = GraphConfig(name="graph-missing-guards")
|
|
context.graph = _fresh_graph(config)
|
|
context.graph.nodes.pop("start", None)
|
|
|
|
|
|
@when("I validate the graph explicitly")
|
|
def step_validate_graph(context):
|
|
context.validation_error = None
|
|
try:
|
|
context.graph._validate_graph() # pylint: disable=protected-access
|
|
except Exception as exc: # pylint: disable=broad-except
|
|
context.validation_error = exc
|
|
|
|
|
|
@then("a guard validation error should be raised")
|
|
def step_assert_guard_validation(context):
|
|
assert isinstance(context.validation_error, ValueError)
|
|
|
|
|
|
@given("a langgraph instance ready to start")
|
|
def step_ready_to_start(context):
|
|
config = GraphConfig(name="graph-start")
|
|
context.graph = _fresh_graph(config)
|
|
context.start_stream = f"__{context.graph.name}_node_start__"
|
|
context.graph.stream_router.send_message = MagicMock()
|
|
|
|
|
|
@when("I start without input then stop and fetch state")
|
|
def step_start_and_stop(context):
|
|
context.graph.start()
|
|
context.graph.stop()
|
|
context.final_state = context.graph.get_state()
|
|
|
|
|
|
@then("it should use default state and update running flags")
|
|
def step_assert_start_stop(context):
|
|
context.graph.stream_router.send_message.assert_called_with(
|
|
context.start_stream, GraphState()
|
|
)
|
|
assert context.graph.is_running is False
|
|
assert isinstance(context.final_state, GraphState)
|
|
|
|
|
|
def _fresh_graph(config: GraphConfig) -> LangGraph:
|
|
"""Construct a LangGraph with isolated scheduler and router."""
|
|
return LangGraph(config=config, stream_router=None, scheduler=None)
|
|
|
|
|
|
@given("a langgraph config with only start and end nodes")
|
|
def step_minimal_config(context):
|
|
context.graph_config = GraphConfig(name="graph-minimal")
|
|
|
|
|
|
@given("asyncio has no running loop")
|
|
def step_no_running_loop(_context):
|
|
try:
|
|
loop = asyncio.get_running_loop()
|
|
if loop.is_running():
|
|
loop.stop()
|
|
except RuntimeError:
|
|
# Expected when no loop is running
|
|
pass
|
|
|
|
|
|
@when("I construct a LangGraph without providing a scheduler")
|
|
def step_construct_graph(context):
|
|
context.graph = _fresh_graph(context.graph_config)
|
|
|
|
|
|
@then("it should create an AsyncIOScheduler and initialize streams")
|
|
def step_assert_scheduler_and_streams(context):
|
|
assert isinstance(context.graph.scheduler, AsyncIOScheduler)
|
|
start_stream = f"__{context.graph.name}_node_start__"
|
|
assert start_stream in context.graph.stream_router.streams
|
|
|
|
|
|
@given("a langgraph instance with a fresh state manager")
|
|
def step_graph_with_state_manager(context):
|
|
config = GraphConfig(name="graph-exec")
|
|
context.graph = _fresh_graph(config)
|
|
context.start_stream = f"__{context.graph.name}_node_start__"
|
|
context.graph.stream_router.send_message = MagicMock()
|
|
|
|
|
|
@when("I execute the graph with input data")
|
|
def step_execute_graph(context):
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
result = loop.run_until_complete(context.graph.execute({"messages": []}))
|
|
loop.close()
|
|
asyncio.set_event_loop(asyncio.new_event_loop())
|
|
context.execute_result = result
|
|
|
|
|
|
@then("the start stream should receive the state and execution returns a GraphState")
|
|
def step_assert_execute(context):
|
|
context.graph.stream_router.send_message.assert_called_once()
|
|
called_state = context.graph.stream_router.send_message.call_args.args[1]
|
|
assert isinstance(called_state, GraphState)
|
|
assert isinstance(context.execute_result, GraphState)
|
|
|
|
|
|
@given("a langgraph instance with a removed start stream")
|
|
def step_graph_missing_start(context):
|
|
config = GraphConfig(name="graph-missing-start")
|
|
context.graph = _fresh_graph(config)
|
|
start_stream = f"__{context.graph.name}_node_start__"
|
|
context.graph.stream_router.streams.pop(start_stream, None)
|
|
|
|
|
|
@when("I call start on the graph")
|
|
def step_call_start(context):
|
|
context.start_error = None
|
|
try:
|
|
context.graph.start()
|
|
except Exception as exc: # pylint: disable=broad-except
|
|
context.start_error = exc
|
|
|
|
|
|
@then("a start stream error should be raised")
|
|
def step_assert_start_error(context):
|
|
assert isinstance(context.start_error, ValueError)
|
|
|
|
|
|
@given("a langgraph config with parallel execution disabled and a cycle")
|
|
def step_cycle_config(context):
|
|
nodes = {
|
|
"a": NodeConfig(name="a", type=NodeType.FUNCTION),
|
|
}
|
|
edges = [Edge(source="start", target="a"), Edge(source="a", target="start")]
|
|
context.graph_config = GraphConfig(
|
|
name="graph-cycle",
|
|
nodes=nodes,
|
|
edges=edges,
|
|
parallel_execution=False,
|
|
)
|
|
|
|
|
|
@when("I construct the LangGraph")
|
|
def step_construct_cycle_graph(context):
|
|
context.graph = _fresh_graph(context.graph_config)
|
|
|
|
|
|
@then("the graph should report cycles and empty parallel groups")
|
|
def step_assert_cycles(context):
|
|
assert context.graph.has_cycles is True
|
|
assert context.graph.parallel_groups == []
|
|
|
|
|
|
@given("a langgraph instance with a patched state manager")
|
|
def step_graph_with_executor(context):
|
|
nodes = {
|
|
"worker": NodeConfig(name="worker", type=NodeType.FUNCTION),
|
|
}
|
|
config = GraphConfig(name="graph-executor", nodes=nodes)
|
|
graph = _fresh_graph(config)
|
|
graph.state_manager.get_state = MagicMock(return_value=GraphState())
|
|
graph.state_manager.update_state = MagicMock()
|
|
graph.nodes["worker"].execute = AsyncMock(return_value={"messages": ["m"]})
|
|
context.executor = graph.stream_router._builtin_execute_node_worker
|
|
context.graph = graph
|
|
|
|
|
|
@when("I invoke the builtin node executor")
|
|
def step_invoke_executor(context):
|
|
result = context.executor(StreamMessage(content="x", metadata={}))
|
|
context.executor_result = result
|
|
|
|
|
|
@then("it should update state and record execution history")
|
|
def step_assert_executor(context):
|
|
context.graph.state_manager.update_state.assert_called_once()
|
|
assert context.graph.execution_history == ["worker"]
|
|
assert isinstance(context.executor_result, StreamMessage)
|