Files
cleveragents-core/features/steps/route_bridge_coverage_steps.py

988 lines
33 KiB
Python

"""Route bridge coverage steps (ported and adapted)."""
from __future__ import annotations
import asyncio
from pathlib import Path
from behave import given, then, when
from rx.scheduler.eventloop import AsyncIOScheduler # type: ignore[attr-defined]
from cleveragents.core.exceptions import ConfigurationError
from cleveragents.langgraph.graph import GraphConfig
from cleveragents.langgraph.nodes import Edge, NodeConfig, NodeType
from cleveragents.langgraph.state import GraphState
from cleveragents.reactive.route import BridgeConfig, RouteConfig, RouteType
from cleveragents.reactive.route_bridge import RouteBridge
from cleveragents.reactive.stream_router import (
ReactiveStreamRouter,
StreamMessage,
)
_DEF_CHECKPOINT_DIR = Path("/tmp/langgraph_checkpoints")
# Reuse a single event loop for async helpers in tests
_TEST_LOOP = asyncio.new_event_loop()
asyncio.set_event_loop(_TEST_LOOP)
def _standard_topological_levels(nodes: dict[str, object]) -> dict[int, set[str]]:
middle_nodes = {name for name in nodes if name not in {"start", "end"}}
return {0: {"start"}, 1: middle_nodes, 2: {"end"}}
def _build_standard_langgraph_stub(corrupt_state_manager: bool = False):
start_cfg = NodeConfig(name="start", type=NodeType.START)
agent_node_cfg = NodeConfig(
name="agent_node", type=NodeType.AGENT, agent="dummy", parallel=False
)
function_node_cfg = NodeConfig(
name="function_node",
type=NodeType.FUNCTION,
function="stub_function",
parallel=False,
)
end_cfg = NodeConfig(name="end", type=NodeType.END)
def _stub_node(cfg):
return type("StubNode", (), {"config": cfg})()
nodes = {
"start": _stub_node(start_cfg),
"agent_node": _stub_node(agent_node_cfg),
"function_node": _stub_node(function_node_cfg),
"end": _stub_node(end_cfg),
}
edges = [
Edge(source="start", target="agent_node"),
Edge(source="agent_node", target="function_node"),
Edge(source="function_node", target="end"),
]
graph_config = GraphConfig(
name="stub_graph",
nodes={name: node.config for name, node in nodes.items()},
edges=edges,
entry_point="start",
checkpointing=True,
checkpoint_dir=_DEF_CHECKPOINT_DIR,
parallel_execution=False,
)
class StubStateManager:
def __init__(self, checkpoint_dir: Path):
self.checkpoint_dir = checkpoint_dir
self.update_count = 0
self.updated_state = None
self.update_state_called = False
self.checkpoint_saved = False
def _save_checkpoint(self): # pylint: disable=protected-access
self.checkpoint_saved = True
return None
def update_state(self, state=None):
self.update_state_called = True
self.updated_state = state or {}
self.update_count += 1
def get_state(self):
return self.updated_state or {}
langgraph = type("StubLangGraph", (), {})()
langgraph.nodes = nodes
langgraph.config = graph_config
langgraph.state_manager = (
None if corrupt_state_manager else StubStateManager(_DEF_CHECKPOINT_DIR)
)
langgraph._topological_levels = lambda: _standard_topological_levels(nodes) # type: ignore[attr-defined]
return langgraph
def _build_stub_langgraph_from_config(
graph_config: GraphConfig, checkpoint_dir: Path | None = None
):
class StubStateManager:
def __init__(self, ckp: Path | None):
self.checkpoint_dir = ckp
self.update_count = 0
self.update_state_called = False
self.updated_state = None
def update_state(self, state=None):
self.update_state_called = True
self.updated_state = state or {}
self.update_count += 1
def get_state(self):
return self.updated_state or {}
def _save_checkpoint(self): # pylint: disable=protected-access
return None
nodes = {
name: type("StubNode", (), {"config": cfg})()
for name, cfg in graph_config.nodes.items()
}
if "start" not in nodes:
nodes["start"] = type(
"StubNode", (), {"config": NodeConfig(name="start", type=NodeType.START)}
)()
if "end" not in nodes:
nodes["end"] = type(
"StubNode", (), {"config": NodeConfig(name="end", type=NodeType.END)}
)()
langgraph = type("StubLangGraph", (), {})()
langgraph.config = graph_config
langgraph.nodes = nodes
langgraph.state_manager = StubStateManager(
checkpoint_dir or graph_config.checkpoint_dir or _DEF_CHECKPOINT_DIR
)
langgraph._topological_levels = lambda: _standard_topological_levels(nodes) # type: ignore[attr-defined]
return langgraph
def _ensure_topological_levels(graph):
if not hasattr(graph, "_topological_levels"):
graph._topological_levels = lambda: _standard_topological_levels( # type: ignore[attr-defined]
getattr(graph, "nodes", {})
)
def _run_async(coro):
try:
loop = asyncio.get_event_loop()
if loop.is_closed():
raise RuntimeError
if loop.is_running():
return asyncio.run(coro)
except Exception: # pylint: disable=broad-except
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
return loop.run_until_complete(coro)
@given("I have a clean test environment for route bridge")
def step_clean_env(context):
context.stream_router = ReactiveStreamRouter()
context.agents = {"agentA": DummyAgent("agentA")}
context.route_bridge = None
context.stream_message = None
context.langgraph = None
context.graph_result = None
context.stream_config_result = None
context.graph_config_result = None
context.route_bridge_error = None
context.state_flattener_called = False
class DummyAgent:
def __init__(self, name: str):
self.name = name
def process_message_sync(self, message, metadata=None):
return f"processed:{message}"
def get_capabilities(self):
return []
@given("I have a route bridge stream router and agents")
def step_stream_router_agents(context):
context.stream_router = ReactiveStreamRouter()
context.agents = {"agentA": DummyAgent("agentA"), "dummy": DummyAgent("dummy")}
for name, agent in context.agents.items():
context.stream_router.register_agent(name, agent)
@given("I have an AsyncIO scheduler")
def step_asyncio_scheduler(context):
try:
loop = asyncio.get_running_loop()
except RuntimeError:
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
context.scheduler = AsyncIOScheduler(loop)
@when("I create the route bridge instance")
def step_create_route_bridge(context):
context.route_bridge = RouteBridge(
context.stream_router,
context.agents,
scheduler=getattr(context, "scheduler", None),
)
@when("I create a RouteBridge instance with scheduler")
def step_create_route_bridge_scheduler(context):
step_create_route_bridge(context)
@then("the route bridge should be initialized correctly")
def step_rb_initialized(context):
assert context.route_bridge is not None
assert hasattr(context.route_bridge, "stream_router")
@then("the route bridge logger should be set up")
def step_logger_set(context):
assert context.route_bridge.logger is not None
@then("route bridge active conversions should be empty")
def step_active_conversions_empty(context):
assert context.route_bridge._active_conversions == {}
@then("the route bridge should be initialized with scheduler")
def step_scheduler_stored(context):
assert context.route_bridge.scheduler is not None
@then("the route bridge scheduler should be stored correctly")
def step_scheduler_correct(context):
assert context.route_bridge.scheduler == getattr(context, "scheduler", None)
@given("I have a RouteBridge instance")
def step_have_route_bridge(context):
if not getattr(context, "route_bridge", None):
step_create_route_bridge(context)
# Reset results for clarity
context.upgrade_result = False
context.downgrade_result = False
context.upgrade_error = None
context.downgrade_error = None
context.route_bridge._active_conversions.clear()
@given("I have a route config without bridge configuration")
def step_route_no_bridge(context):
context.route_config = RouteConfig(
name="r1",
type=RouteType.GRAPH,
nodes={
"agent_node": {"type": "agent", "agent": "dummy", "parallel": False},
"function_node": {
"type": "function",
"function": "stub_function",
"parallel": False,
},
},
edges=[
{"source": "start", "target": "agent_node"},
{"source": "agent_node", "target": "function_node"},
{"source": "function_node", "target": "end"},
],
entry_point="start",
checkpointing=True,
checkpoint_dir=str(_DEF_CHECKPOINT_DIR),
)
context.downgrade_result = False
context.upgrade_result = False
@given("I have a graph route config with bridge configuration")
def step_graph_route_with_bridge(context):
context.route_config = RouteConfig(
name="g1",
type=RouteType.GRAPH,
nodes={
"agent_node": {"type": "agent", "agent": "dummy", "parallel": False},
"function_node": {
"type": "function",
"function": "stub_function",
"parallel": False,
},
},
edges=[
{"source": "start", "target": "agent_node"},
{"source": "agent_node", "target": "function_node"},
{"source": "function_node", "target": "end"},
],
entry_point="start",
bridge=BridgeConfig(),
checkpointing=True,
checkpoint_dir=str(_DEF_CHECKPOINT_DIR),
parallel_execution=False,
)
@given("I have a stream route config with needs_state upgrade condition")
def step_stream_route_needs_state(context):
context.route_config = RouteConfig(
name="needs_state",
type=RouteType.STREAM,
operators=[{"type": "map", "params": {"agent": "dummy"}}],
bridge=BridgeConfig(upgrade_conditions={"needs_state": True}),
)
@given("I have a stream route config with message_count upgrade condition")
def step_stream_route_message_count(context):
context.route_config = RouteConfig(
name="message_count",
type=RouteType.STREAM,
operators=[{"type": "map", "params": {"agent": "dummy"}}],
bridge=BridgeConfig(upgrade_conditions={"message_count": 5}),
)
@given("I have a stream route config with complexity_threshold upgrade condition")
def step_stream_route_complexity_threshold(context):
context.route_config = RouteConfig(
name="complexity_threshold",
type=RouteType.STREAM,
operators=[
{"type": "map", "params": {"agent": "dummy"}},
{"type": "map", "params": {"function": "f"}},
],
bridge=BridgeConfig(upgrade_conditions={"complexity_threshold": 6}),
)
@given("I have a stream route config with custom_predicate upgrade condition")
def step_stream_route_custom_predicate(context):
context.route_config = RouteConfig(
name="custom_predicate",
type=RouteType.STREAM,
operators=[{"type": "map", "params": {"agent": "dummy"}}],
bridge=BridgeConfig(upgrade_conditions={"custom_predicate": None}),
)
@given("I have a stream route config with non-callable custom_predicate")
def step_stream_route_non_callable_predicate(context):
context.route_config = RouteConfig(
name="non_callable_predicate",
type=RouteType.STREAM,
operators=[{"type": "map", "params": {"agent": "dummy"}}],
bridge=BridgeConfig(upgrade_conditions={"custom_predicate": "not_callable"}),
)
@given("I have a stream message")
def step_stream_message(context):
context.stream_message = StreamMessage(content="data", metadata={})
@given("I have a stream message with requires_state metadata")
def step_stream_message_requires_state(context):
context.stream_message = StreamMessage(
content="data", metadata={"requires_state": True}
)
@given("I have a graph route config for downgrade")
def step_graph_route_for_downgrade(context):
context.route_config = RouteConfig(
name="graph_to_stream",
type=RouteType.GRAPH,
nodes={
"agent_node": {"type": "agent", "agent": "dummy", "parallel": False},
"function_node": {
"type": "function",
"function": "stub_function",
"parallel": False,
},
},
edges=[
{"source": "start", "target": "agent_node"},
{"source": "agent_node", "target": "function_node"},
{"source": "function_node", "target": "end"},
],
entry_point="start",
checkpointing=True,
checkpoint_dir=str(_DEF_CHECKPOINT_DIR),
parallel_execution=False,
bridge=BridgeConfig(),
)
@given("I have a LangGraph instance")
def step_have_langgraph_instance(context):
context.langgraph = _build_standard_langgraph_stub()
_ensure_topological_levels(context.langgraph)
context.route_bridge._active_conversions[context.route_config.name] = {
"type": "graph",
"instance": context.langgraph,
"original_config": context.route_config,
}
@given("I have a graph state")
def step_graph_state(context):
context.graph_state = GraphState()
context.graph_state.metadata = {}
context.graph_state.messages = []
@given("I have a graph route config with idle_time downgrade condition")
def step_graph_route_idle_time(context):
bridge_config = BridgeConfig(downgrade_conditions={"idle_time": 5.0})
context.route_config = RouteConfig(
name="graph_idle",
type=RouteType.GRAPH,
nodes={"n1": {"type": "agent", "agent": "dummy"}},
edges=[{"source": "start", "target": "n1"}],
bridge=bridge_config,
)
@given("I have a graph state that has been idle too long")
def step_graph_state_idle(context):
context.graph_state = GraphState()
current_time = asyncio.get_event_loop().time()
context.graph_state.metadata = {"last_updated": current_time - 10.0}
context.graph_state.messages = []
@given("I have a graph route config with state_size downgrade condition")
def step_graph_route_state_size(context):
bridge_config = BridgeConfig(downgrade_conditions={"state_size": 1})
context.route_config = RouteConfig(
name="graph_state_size",
type=RouteType.GRAPH,
nodes={"n1": {"type": "agent", "agent": "dummy"}},
edges=[{"source": "start", "target": "n1"}],
bridge=bridge_config,
)
@given("I have a graph state with minimal messages")
def step_graph_state_minimal(context):
context.graph_state = GraphState()
context.graph_state.metadata = {}
context.graph_state.messages = [{"content": "only_one"}]
@given("I have a graph route config with no_conditionals_used downgrade condition")
def step_graph_route_no_conditionals(context):
bridge_config = BridgeConfig(downgrade_conditions={"no_conditionals_used": True})
context.route_config = RouteConfig(
name="graph_no_conditionals",
type=RouteType.GRAPH,
nodes={"n1": {"type": "agent", "agent": "dummy"}},
edges=[{"source": "start", "target": "n1"}],
bridge=bridge_config,
)
@given("I have a graph state with no conditional edges used")
def step_graph_state_no_conditionals(context):
context.graph_state = GraphState()
context.graph_state.metadata = {"conditional_edges_used": False}
context.graph_state.messages = []
@given("I have a stream route config with bridge configuration")
def step_stream_route_with_bridge_config(context):
bridge_config = BridgeConfig(
upgrade_conditions={"needs_state": True},
downgrade_conditions={"idle_time": 30.0},
)
context.route_config = RouteConfig(
name="stream_with_bridge",
type=RouteType.STREAM,
operators=[{"type": "map", "params": {"agent": "dummy"}}],
bridge=bridge_config,
)
@when("I check downgrade conditions")
def step_check_downgrade(context):
if not getattr(context, "graph_state", None):
step_graph_state(context)
try:
context.downgrade_result = _run_async(
context.route_bridge.check_downgrade_conditions(
context.route_config, context.graph_state
)
)
except Exception as exc: # pylint: disable=broad-except
context.downgrade_result = False
context.downgrade_error = exc
@then("the downgrade result should be True")
def step_downgrade_result_true(context):
assert context.downgrade_result is True
@then("the downgrade result should be False")
def step_downgrade_result_false(context):
assert context.downgrade_result is False
@given("I have processed enough messages to trigger upgrade")
def step_processed_enough_messages(context):
dummy_instance = type("DummyConversion", (), {"update_count": 5})()
context.route_bridge._active_conversions[context.route_config.name] = {
"type": "stream",
"instance": dummy_instance,
"original_config": context.route_config,
}
@given("I have a complex route configuration")
def step_complex_route_configuration(context):
# Ensure the route has enough operators to exceed the threshold
if (
not getattr(context, "route_config", None)
or context.route_config.type != RouteType.STREAM
):
context.route_config = RouteConfig(
name="complex_route",
type=RouteType.STREAM,
operators=[
{"type": "map", "params": {"agent": "dummy"}},
{"type": "map", "params": {"function": "f1"}},
{"type": "map", "params": {"function": "f2"}},
{"type": "map", "params": {"function": "f3"}},
{"type": "map", "params": {"function": "f4"}},
],
bridge=BridgeConfig(upgrade_conditions={"complexity_threshold": 6}),
)
else:
# Enrich existing operators to make the route more complex
context.route_config.operators.extend(
[
{"type": "map", "params": {"function": "extra_f1"}},
{"type": "map", "params": {"function": "extra_f2"}},
{"type": "map", "params": {"function": "extra_f3"}},
]
)
@given("I seed a message count for the route")
def step_seed_message_count(context):
if not getattr(context, "route_config", None):
context.route_config = RouteConfig(
name="message_count_route",
type=RouteType.STREAM,
operators=[{"type": "map", "params": {"agent": "dummy"}}],
bridge=BridgeConfig(upgrade_conditions={"message_count": 5}),
)
state_manager = type("DummyStateManager", (), {"update_count": 7})()
dummy_graph = type("DummyGraph", (), {"state_manager": state_manager})()
context.route_bridge._active_conversions[context.route_config.name] = {
"type": "graph",
"instance": dummy_graph,
"original_config": context.route_config,
}
@when("I get message count for a route")
def step_get_message_count(context):
if not getattr(context, "route_config", None):
step_seed_message_count(context)
context.message_count = context.route_bridge._get_message_count(
context.route_config.name
)
@when("I check upgrade conditions")
def step_check_upgrade(context):
if not getattr(context, "stream_message", None):
context.stream_message = StreamMessage(content="data", metadata={})
try:
context.upgrade_result = _run_async(
context.route_bridge.check_upgrade_conditions(
context.route_config, context.stream_message
)
)
except Exception as exc: # pylint: disable=broad-except
context.upgrade_result = False
context.upgrade_error = exc
@when("I check upgrade conditions synchronously")
def step_check_upgrade_sync(context):
if not getattr(context, "stream_message", None):
context.stream_message = StreamMessage(content="data", metadata={})
coro = context.route_bridge.check_upgrade_conditions(
context.route_config, context.stream_message
)
context.upgrade_result = _run_async(coro)
@then("the route bridge upgrade result should be True")
def step_upgrade_result_true(context):
assert context.upgrade_result is True
@then("the route bridge upgrade result should be False")
def step_upgrade_result_false(context):
assert context.upgrade_result is False
@then("the route bridge message count should be at least 5")
def step_message_count_at_least(context):
assert context.route_bridge._get_message_count(context.route_config.name) >= 5
@then("the route bridge message count should be returned")
def step_message_count_returned(context):
assert context.message_count is not None
@then("the route bridge message count result should be a valid integer")
def step_message_count_valid_integer(context):
assert isinstance(context.message_count, int)
@then("route bridge error handling should occur")
def step_error_handling(context):
assert context.upgrade_error is None or isinstance(context.upgrade_error, Exception)
@then("the route bridge system should remain stable")
def step_system_stable(context):
assert context.route_bridge is not None
@given("I have an invalid graph route config")
def step_invalid_graph_route_config(context):
# Build an invalid placeholder but catch validation so test can proceed
try:
context.route_config = RouteConfig(
name="bad_graph",
type=RouteType.GRAPH,
nodes={},
edges=[],
)
except ConfigurationError as exc:
context.route_config = None
context.downgrade_error = exc
@given("I have a corrupted LangGraph instance")
def step_corrupted_langgraph(context):
context.langgraph = _build_standard_langgraph_stub(corrupt_state_manager=True)
@when("I attempt to downgrade graph to stream")
def step_attempt_downgrade(context):
if context.route_config is None:
context.downgrade_result = False
context.downgrade_error = context.downgrade_error or ConfigurationError(
"Invalid graph route"
)
return
try:
context.downgrade_result = asyncio.get_event_loop().run_until_complete(
context.route_bridge.downgrade_graph_to_stream(
context.route_config, context.langgraph
)
)
context.downgrade_error = None
except ConfigurationError as exc:
context.downgrade_result = False
context.downgrade_error = exc
except Exception as exc: # pylint: disable=broad-except
context.downgrade_error = exc
@then("And the system should remain stable")
def step_system_stable_downgrade(context):
assert context.route_bridge is not None
@given("I have a route config that supports both upgrade and downgrade")
def step_route_upgrade_downgrade(context):
context.route_config = RouteConfig(
name="cycle",
type=RouteType.STREAM,
operators=[{"type": "map", "params": {"agent": "dummy"}}],
bridge=BridgeConfig(
upgrade_conditions={"needs_state": True},
downgrade_conditions={"state_size": 0},
),
)
context.stream_message = StreamMessage(
content="x", metadata={"requires_state": True}
)
@when("I perform an upgrade and then downgrade cycle")
def step_cycle_upgrade_downgrade(context):
context.langgraph = asyncio.get_event_loop().run_until_complete(
context.route_bridge.upgrade_stream_to_graph(
context.route_config, context.stream_message
)
)
context.route_bridge._active_conversions[context.route_config.name] = {
"type": "graph",
"instance": context.langgraph,
"original_config": context.route_config,
}
_ensure_topological_levels(context.langgraph)
# Convert Node objects to NodeConfig dict
node_configs = {}
for name, node in context.langgraph.nodes.items():
if hasattr(node, "config") and node.config:
node_configs[name] = node.config
else:
# Create a minimal NodeConfig for nodes without config
node_configs[name] = NodeConfig(
name=name, type=getattr(node, "type", NodeType.FUNCTION)
)
context.stream_config = asyncio.get_event_loop().run_until_complete(
context.route_bridge.downgrade_graph_to_stream(
RouteConfig.from_graph_config(
GraphConfig(
name="cycle_g",
nodes=node_configs,
edges=context.langgraph.config.edges,
entry_point="start",
)
),
context.langgraph,
)
)
@then("the cycle should complete successfully")
def step_cycle_success(context):
assert context.stream_config is not None
@then("the final configuration should be consistent")
def step_final_config_consistent(context):
assert context.stream_config.name
@then("active conversions should be tracked correctly")
def step_active_conversions_tracked(context):
assert context.route_bridge._active_conversions
# Additional coverage steps
class _FailingRegistry:
def list_actors(self):
raise RuntimeError("registry failure")
@given("I have a failing agent registry for route bridge")
def step_failing_registry(context):
context.stream_router = ReactiveStreamRouter()
context.agents = _FailingRegistry()
@then("the route bridge agents should be empty")
def step_agents_empty(context):
assert context.route_bridge.agents == {}
@given("I have a pending message count below threshold")
def step_message_count_below_threshold(context):
dummy_instance = type("DummyConversion", (), {"update_count": 2})()
context.route_bridge._active_conversions[context.route_config.name] = {
"type": "stream",
"instance": dummy_instance,
"original_config": context.route_config,
}
@given("I have a simple stream route config with high complexity threshold")
def step_simple_route_high_complexity(context):
context.route_config = RouteConfig(
name="simple_high_threshold",
type=RouteType.STREAM,
operators=[{"type": "map", "params": {"agent": "dummy"}}],
bridge=BridgeConfig(upgrade_conditions={"complexity_threshold": 999}),
)
@given("I have a stream route config with callable custom_predicate")
def step_route_callable_predicate(context):
def _predicate(message, _cfg):
return True
bridge_config = BridgeConfig(upgrade_conditions={"custom_predicate": None})
object.__setattr__(
bridge_config, "upgrade_conditions", {"custom_predicate": _predicate}
)
context.route_config = RouteConfig(
name="callable_predicate",
type=RouteType.STREAM,
operators=[{"type": "map", "params": {"agent": "dummy"}}],
bridge=bridge_config,
)
@given("I have a stream route config with state extractor and subscriptions")
def step_route_with_state_extractor(context):
def _extractor(_msg):
return {"metadata": {"extracted": True}}
bridge_config = BridgeConfig(preserve_subscriptions=True)
object.__setattr__(bridge_config, "state_extractor", _extractor)
context.route_config = RouteConfig(
name="extractor_route",
type=RouteType.STREAM,
operators=[{"type": "map", "params": {"agent": "dummy"}}],
subscriptions=["a", "b"],
bridge=bridge_config,
)
@when("I create a graph config from a stream without agents")
def step_create_graph_config_without_agents(context):
context.route_config = RouteConfig(
name="function_only_stream",
type=RouteType.STREAM,
operators=[{"type": "map", "params": {"function": "noop"}}],
)
context.graph_config_result = context.route_bridge._create_graph_from_stream(
context.route_config
)
@given("I seed a fallback update_count conversion")
def step_seed_fallback_update_count(context):
context.route_config = RouteConfig(
name="update_count_route",
type=RouteType.STREAM,
operators=[{"type": "map", "params": {"agent": "dummy"}}],
bridge=BridgeConfig(upgrade_conditions={"message_count": 5}),
)
instance = type("DummyConversion", (), {"update_count": 3})()
context.route_bridge._active_conversions[context.route_config.name] = {
"type": "graph",
"instance": instance,
"original_config": context.route_config,
}
@when("I upgrade the stream route to graph")
def step_upgrade_stream_with_extractor(context):
if not getattr(context, "stream_message", None):
context.stream_message = StreamMessage(content="data", metadata={})
context.graph_result = _run_async(
context.route_bridge.upgrade_stream_to_graph(
context.route_config, context.stream_message
)
)
@given("I have a graph route config with flattener and checkpointing")
def step_graph_route_with_flattener(context):
def _flattener(state):
context.state_flattener_called = True
return state
bridge_config = BridgeConfig(preserve_checkpointing=True)
object.__setattr__(bridge_config, "state_flattener", _flattener)
context.route_config = RouteConfig(
name="graph_flat",
type=RouteType.GRAPH,
nodes={"n1": {"type": "agent", "agent": "dummy"}},
edges=[{"source": "start", "target": "n1"}],
bridge=bridge_config,
checkpointing=True,
checkpoint_dir=str(_DEF_CHECKPOINT_DIR),
parallel_execution=False,
)
@given("I have a LangGraph instance with checkpointing")
def step_langgraph_with_checkpointing(context):
context.langgraph = _build_standard_langgraph_stub()
# ensure checkpoint_dir persists for test assertions
if hasattr(context.langgraph, "state_manager"):
context.langgraph.state_manager.checkpoint_dir = _DEF_CHECKPOINT_DIR
@when("I downgrade the graph with flattener")
def step_downgrade_with_flattener(context):
context.stream_config_result = _run_async(
context.route_bridge.downgrade_graph_to_stream(
context.route_config, context.langgraph
)
)
@then("the state flattener should run")
def step_state_flattener_run(context):
assert context.state_flattener_called is True
@then("the checkpoint should be saved")
def step_checkpoint_saved_flag(context):
assert getattr(context.langgraph.state_manager, "checkpoint_saved", False) is True
@then("the graph initial state should include extracted data")
def step_graph_state_includes_extracted(context):
state_mgr = getattr(context.graph_result, "state_manager", None)
assert state_mgr is not None
state = state_mgr.get_state()
assert isinstance(state, GraphState)
assert state.metadata.get("extracted") is True
@given("I have a graph state with multiple messages")
def step_graph_state_multiple_messages(context):
context.graph_state = GraphState()
context.graph_state.metadata = {}
context.graph_state.messages = [
{"content": "a"},
{"content": "b"},
{"content": "c"},
]
@given("I have a graph route config with idle_time not exceeded")
def step_graph_route_idle_time_not_exceeded(context):
bridge_config = BridgeConfig(downgrade_conditions={"idle_time": 1000.0})
context.route_config = RouteConfig(
name="graph_idle_safe",
type=RouteType.GRAPH,
nodes={"n1": {"type": "agent", "agent": "dummy"}},
edges=[{"source": "start", "target": "n1"}],
bridge=bridge_config,
)
context.graph_state = GraphState()
context.graph_state.metadata = {"last_updated": _TEST_LOOP.time()}
context.graph_state.messages = []
@then("the route bridge fallback message count should be returned")
def step_fallback_message_count(context):
assert context.route_bridge._get_message_count(context.route_config.name) == 3
@given("I have a graph route config with large state size")
def step_graph_route_large_state(context):
bridge_config = BridgeConfig(downgrade_conditions={"state_size": 1})
context.route_config = RouteConfig(
name="graph_state_large",
type=RouteType.GRAPH,
nodes={"n1": {"type": "agent", "agent": "dummy"}},
edges=[{"source": "start", "target": "n1"}],
bridge=bridge_config,
)
context.graph_state = GraphState()
context.graph_state.metadata = {}
context.graph_state.messages = [
{"content": "a"},
{"content": "b"},
{"content": "c"},
]
@then("the resulting graph should use function nodes")
def step_result_graph_function_nodes(context):
node_cfg = context.graph_config_result.nodes.get("op_0_map")
assert node_cfg is not None
assert node_cfg.type == NodeType.FUNCTION