290 lines
9.4 KiB
Python
290 lines
9.4 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import rx
|
|
from behave import given, then, when
|
|
|
|
from cleveragents.langgraph.bridge import RxPyLangGraphBridge
|
|
from cleveragents.langgraph.state import GraphState
|
|
from cleveragents.reactive.stream_router import ReactiveStreamRouter, StreamMessage
|
|
|
|
|
|
def _make_graph_state(
|
|
messages: list | None = None, data: dict | None = None
|
|
) -> GraphState:
|
|
state = MagicMock(spec=GraphState)
|
|
state.messages = messages or []
|
|
state.to_dict = MagicMock(return_value=data or {"messages": state.messages})
|
|
return state
|
|
|
|
|
|
@given("I set up a langgraph bridge test harness")
|
|
def step_impl(context):
|
|
scheduler = MagicMock()
|
|
router = ReactiveStreamRouter(scheduler=scheduler)
|
|
router.agents = {"agent": MagicMock()}
|
|
context.router = router
|
|
context.bridge = RxPyLangGraphBridge(router)
|
|
context.results = []
|
|
|
|
|
|
@when("I create a bridge with mocked graph configuration")
|
|
def step_impl(context):
|
|
config = {
|
|
"name": "test_graph",
|
|
"entry_point": "start",
|
|
"nodes": {
|
|
"worker": {
|
|
"type": "function",
|
|
"function": "do_work",
|
|
"metadata": {"priority": "high"},
|
|
}
|
|
},
|
|
"edges": [{"source": "start", "target": "worker"}],
|
|
}
|
|
|
|
with patch("cleveragents.langgraph.bridge.LangGraph") as mock_langgraph:
|
|
mock_graph = MagicMock()
|
|
mock_graph.name = "test_graph"
|
|
mock_graph.nodes = {}
|
|
mock_graph.state_manager = MagicMock()
|
|
mock_graph.state_manager.get_state = MagicMock()
|
|
mock_graph.state_manager.update_state = MagicMock()
|
|
mock_graph.state_manager._save_checkpoint = MagicMock()
|
|
mock_graph.get_execution_history = MagicMock(return_value=["step"])
|
|
mock_langgraph.return_value = mock_graph
|
|
|
|
graph = context.bridge.create_graph_from_config(config)
|
|
context.graph = graph
|
|
context.graph_mock = mock_graph
|
|
context.results.append("create_graph_from_config")
|
|
|
|
# operator registration & helper passthrough
|
|
context.results.append(
|
|
"operators_registered"
|
|
if hasattr(context.bridge, "_operator_graph_execute")
|
|
else "operators_missing"
|
|
)
|
|
context.results.append(
|
|
"run_async_safely"
|
|
if context.bridge._run_async_safely("x") == "x"
|
|
else "run_async_failed"
|
|
)
|
|
|
|
# cleanup_tasks path with pending task
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
|
|
async def pending():
|
|
await asyncio.sleep(0.01)
|
|
|
|
task = loop.create_task(pending())
|
|
context.bridge._active_tasks.add(task)
|
|
context.bridge.cleanup_tasks()
|
|
loop.run_until_complete(asyncio.sleep(0))
|
|
context.results.append(
|
|
"cleanup_tasks_cancelled" if task.cancelled() else "cleanup_tasks_failed"
|
|
)
|
|
loop.close()
|
|
asyncio.set_event_loop(asyncio.new_event_loop())
|
|
|
|
|
|
@when("I exercise bridge graph and stream creation")
|
|
def step_impl(context):
|
|
stream_config = context.bridge.create_graph_stream("test_graph")
|
|
assert stream_config.name == "graph_test_graph"
|
|
context.results.append("create_graph_stream")
|
|
try:
|
|
context.bridge.create_graph_stream("missing")
|
|
except ValueError:
|
|
context.results.append("create_graph_stream_error")
|
|
|
|
|
|
@when("I exercise graph executor branches")
|
|
def step_impl(context):
|
|
bridge = context.bridge
|
|
graph = context.graph_mock
|
|
|
|
async def exec_with_messages(data):
|
|
return _make_graph_state(messages=[{"content": "ok"}])
|
|
|
|
async def exec_without_messages(data):
|
|
return _make_graph_state(messages=[], data={"data": "empty"})
|
|
|
|
graph.execute = exec_with_messages
|
|
|
|
executor = bridge._create_graph_executor({"graph": "test_graph"})
|
|
loop = asyncio.get_event_loop()
|
|
|
|
def _run_operator(message):
|
|
results = []
|
|
rx.just(message).pipe(executor).subscribe(lambda x: results.append(x))
|
|
loop.run_until_complete(asyncio.sleep(0.01))
|
|
return results
|
|
|
|
_run_operator(StreamMessage(content="hello", metadata={}))
|
|
context.results.append("graph_executor_string")
|
|
|
|
_run_operator(
|
|
StreamMessage(
|
|
content={"messages": [{"role": "user", "content": "hi"}]}, metadata={}
|
|
)
|
|
)
|
|
context.results.append("graph_executor_dict")
|
|
|
|
_run_operator(StreamMessage(content=123, metadata={"meta": True}))
|
|
context.results.append("graph_executor_other")
|
|
|
|
graph.execute = exec_without_messages
|
|
_run_operator(StreamMessage(content="empty", metadata={}))
|
|
context.results.append("graph_executor_empty")
|
|
|
|
try:
|
|
bridge._create_graph_executor({"graph": "bad"})
|
|
except ValueError:
|
|
context.results.append("graph_executor_error")
|
|
|
|
|
|
@when("I exercise state utilities and node operator")
|
|
def step_impl(context):
|
|
bridge = context.bridge
|
|
graph = context.graph_mock
|
|
|
|
updater = bridge._create_state_updater({"graph": "test_graph"})
|
|
rx.just(StreamMessage(content={"value": 1}, metadata={"a": 1})).pipe(
|
|
updater
|
|
).subscribe(lambda _: None)
|
|
rx.just(StreamMessage(content="text", metadata={})).pipe(updater).subscribe(
|
|
lambda _: None
|
|
)
|
|
context.results.append("state_updater")
|
|
|
|
try:
|
|
bridge._create_state_updater({"graph": "bad"})
|
|
except ValueError:
|
|
context.results.append("state_updater_error")
|
|
|
|
graph.state_manager.checkpoint_dir = True
|
|
checkpointer = bridge._create_state_checkpointer({"graph": "test_graph"})
|
|
rx.just(StreamMessage(content="x", metadata={"cp": True})).pipe(
|
|
checkpointer
|
|
).subscribe(lambda _: None)
|
|
graph.state_manager._save_checkpoint.assert_called()
|
|
context.results.append("state_checkpointer")
|
|
|
|
try:
|
|
bridge._create_state_checkpointer({"graph": "bad"})
|
|
except ValueError:
|
|
context.results.append("state_checkpointer_error")
|
|
|
|
node = MagicMock()
|
|
node.execute = AsyncMock(return_value={"messages": [{"content": "node"}]})
|
|
graph.nodes["worker"] = node
|
|
node_op = bridge._create_node_operator({"graph": "test_graph", "node": "worker"})
|
|
rx.just(StreamMessage(content="n", metadata={})).pipe(node_op).subscribe(
|
|
lambda _: None
|
|
)
|
|
context.results.append("node_operator")
|
|
|
|
node.execute = AsyncMock(return_value={"data": "nomsg"})
|
|
rx.just(StreamMessage(content="n", metadata={})).pipe(node_op).subscribe(
|
|
lambda _: None
|
|
)
|
|
context.results.append("node_operator_nomsg")
|
|
|
|
try:
|
|
bridge._create_node_operator({"graph": "bad", "node": "worker"})
|
|
except ValueError:
|
|
context.results.append("node_operator_bad_graph")
|
|
|
|
try:
|
|
bridge._create_node_operator({"graph": "test_graph", "node": "missing"})
|
|
except ValueError:
|
|
context.results.append("node_operator_bad_node")
|
|
|
|
|
|
@when("I exercise conditional routing helpers")
|
|
def step_impl(context):
|
|
bridge = context.bridge
|
|
rules = [
|
|
{"condition": {"equals": "match"}, "target": "t1"},
|
|
{"condition": {"field": "kind", "value": "x"}, "target": "t2"},
|
|
]
|
|
router = bridge._create_conditional_router({"rules": rules})
|
|
|
|
msg = StreamMessage(content="match", metadata={})
|
|
result_list = []
|
|
rx.just(msg).pipe(router).subscribe(lambda x: result_list.append(x))
|
|
context.results.append("conditional_equals")
|
|
|
|
msg2 = StreamMessage(content={"kind": "x"}, metadata={})
|
|
result_list2 = []
|
|
rx.just(msg2).pipe(router).subscribe(lambda x: result_list2.append(x))
|
|
context.results.append("conditional_field")
|
|
|
|
msg3 = StreamMessage(content="none", metadata={})
|
|
result_list3 = []
|
|
rx.just(msg3).pipe(router).subscribe(lambda x: result_list3.append(x))
|
|
context.results.append("conditional_default")
|
|
|
|
assert bridge._evaluate_condition(
|
|
StreamMessage(content="value", metadata={}), {"equals": "value"}
|
|
)
|
|
assert not bridge._evaluate_condition(
|
|
StreamMessage(content={"k": "v"}, metadata={}), {"field": "k", "value": "x"}
|
|
)
|
|
context.results.append("evaluate_condition")
|
|
|
|
|
|
@when("I exercise router node helpers and cleanup")
|
|
def step_impl(context):
|
|
bridge = context.bridge
|
|
|
|
rules = [{"condition": {"equals": "yes"}, "target": "node_a"}]
|
|
node_cfg = bridge.create_message_router_node(rules)
|
|
assert node_cfg.metadata["rules"] == rules
|
|
context.results.append("message_router_node")
|
|
|
|
edges = bridge.create_conditional_edges(rules)
|
|
assert len(edges) == 2 # start->target and start->end
|
|
context.results.append("conditional_edges")
|
|
|
|
bridge.cleanup_tasks()
|
|
context.results.append("cleanup_tasks_empty")
|
|
|
|
|
|
@then("bridge coverage steps should complete successfully")
|
|
def step_impl(context):
|
|
expected = {
|
|
"create_graph_from_config",
|
|
"operators_registered",
|
|
"run_async_safely",
|
|
"cleanup_tasks_cancelled",
|
|
"create_graph_stream",
|
|
"create_graph_stream_error",
|
|
"graph_executor_string",
|
|
"graph_executor_dict",
|
|
"graph_executor_other",
|
|
"graph_executor_empty",
|
|
"graph_executor_error",
|
|
"state_updater",
|
|
"state_updater_error",
|
|
"state_checkpointer",
|
|
"state_checkpointer_error",
|
|
"node_operator",
|
|
"node_operator_nomsg",
|
|
"node_operator_bad_graph",
|
|
"node_operator_bad_node",
|
|
"conditional_equals",
|
|
"conditional_field",
|
|
"conditional_default",
|
|
"evaluate_condition",
|
|
"message_router_node",
|
|
"conditional_edges",
|
|
"cleanup_tasks_empty",
|
|
}
|
|
missing = expected - set(context.results)
|
|
assert not missing, f"Missing coverage steps: {missing}"
|