cleveragents-core/features/steps/langgraph_bridge_steps.py

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"})
    try:
        loop = asyncio.get_running_loop()
    except RuntimeError:
        loop = asyncio.new_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")

    # Close the fresh event loop to prevent resource leak
    with __import__("contextlib").suppress(RuntimeError):
        loop.close()  # Safe: get_running_loop catches if another loop already exists


@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}"