100 lines
3.2 KiB
Python
100 lines
3.2 KiB
Python
from collections import deque
|
|
from typing import Any
|
|
|
|
from behave import given, then, when
|
|
|
|
from cleveragents.langgraph.nodes import Edge, NodeConfig, NodeType
|
|
from cleveragents.langgraph.pure_graph import PureGraph
|
|
|
|
|
|
def _build_graph(name: str, node_defs: list[tuple[str, str | None]]):
|
|
nodes: dict[str, NodeConfig] = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
"end": NodeConfig(name="end", type=NodeType.END),
|
|
}
|
|
edges: list[Edge] = []
|
|
previous = "start"
|
|
for node_name, function_name in node_defs:
|
|
nodes[node_name] = NodeConfig(
|
|
name=node_name, type=NodeType.FUNCTION, function=function_name
|
|
)
|
|
edges.append(Edge(source=previous, target=node_name))
|
|
previous = node_name
|
|
edges.append(Edge(source=previous, target="end"))
|
|
return PureGraph(name=name, nodes=nodes, edges=edges)
|
|
|
|
|
|
@given('a pure graph with nodes "{first}" and "{second}"')
|
|
def step_graph_with_two_nodes(context, first: str, second: str):
|
|
context.graph = _build_graph("pure-two", [(first, None), (second, None)])
|
|
|
|
|
|
@when("I list its topological order")
|
|
def step_list_topological_order(context):
|
|
context.topo_order = context.graph.topological_order()
|
|
|
|
|
|
@then("the order should be start {first} {second} end")
|
|
def step_assert_topological_order(context, first: str, second: str):
|
|
assert context.topo_order == ["start", first, second, "end"]
|
|
|
|
|
|
@given('a pure graph with function nodes "double" and "increment"')
|
|
def step_graph_with_functions(context):
|
|
context.call_sequence = deque()
|
|
|
|
def double(x: Any) -> Any:
|
|
context.call_sequence.append("double")
|
|
return x * 2
|
|
|
|
def increment(x: Any) -> Any:
|
|
context.call_sequence.append("increment")
|
|
return x + 1
|
|
|
|
context.fn_registry = {
|
|
"double": double,
|
|
"increment": increment,
|
|
}
|
|
context.graph = _build_graph(
|
|
"pure-fns", [("double", "double"), ("increment", "increment")]
|
|
)
|
|
|
|
|
|
@given("registered functions that double then increment")
|
|
def step_registered_functions(_context):
|
|
# Functions already registered in previous step
|
|
return
|
|
|
|
|
|
@when("I execute the pure graph starting with {initial:d}")
|
|
def step_execute_graph(context, initial: int):
|
|
context.exec_result = context.graph.execute(context.fn_registry, initial=initial)
|
|
|
|
|
|
@then("the final result should be {expected:d}")
|
|
def step_assert_final_result(context, expected: int):
|
|
assert context.exec_result == expected
|
|
|
|
|
|
@then("the functions should run in declaration order")
|
|
def step_assert_call_order(context):
|
|
assert list(context.call_sequence) == ["double", "increment"]
|
|
|
|
|
|
@given("a pure graph with a missing function node")
|
|
def step_graph_with_missing_function(context):
|
|
context.graph = _build_graph("pure-missing", [("skipper", "missing_fn")])
|
|
context.fn_registry = {}
|
|
|
|
|
|
@then("the result should remain {expected:d}")
|
|
def step_assert_result_unchanged(context, expected: int):
|
|
assert context.exec_result == expected
|
|
|
|
|
|
@given('a pure graph with non-functional nodes "{first}" and "{second}"')
|
|
def step_graph_with_non_functional_nodes(context, first: str, second: str):
|
|
# Build graph without assigning functions and keep registry empty
|
|
context.graph = _build_graph("pure-inert", [(first, None), (second, None)])
|
|
context.fn_registry = {}
|