forked from cleveragents/cleveragents-core
882 lines
30 KiB
Python
882 lines
30 KiB
Python
"""
|
|
Unit tests for langgraph/bridge.py module
|
|
|
|
Tests the bridge between RxPy streams and LangGraph execution.
|
|
"""
|
|
|
|
import pytest
|
|
import asyncio
|
|
from unittest.mock import Mock, AsyncMock
|
|
|
|
import rx
|
|
from rx.subject import Subject
|
|
|
|
from cleveragents.langgraph.bridge import RxPyLangGraphBridge
|
|
from cleveragents.langgraph.graph import LangGraph
|
|
from cleveragents.langgraph.state import GraphState
|
|
from cleveragents.reactive.stream_router import ReactiveStreamRouter, StreamConfig, StreamMessage, StreamType
|
|
from rx.scheduler.eventloop import AsyncIOScheduler
|
|
|
|
|
|
@pytest.fixture
|
|
def scheduler():
|
|
"""Fixture for AsyncIO scheduler.
|
|
|
|
Note: This fixture works for both sync and async tests. It attempts to use
|
|
the running event loop (for async test contexts), or creates a new one if
|
|
needed (for sync test contexts). This avoids the Python 3.12+ deprecation
|
|
warning for asyncio.get_event_loop() when no loop exists.
|
|
"""
|
|
try:
|
|
# Try to get the running loop (for async test contexts)
|
|
loop = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
# No running loop - create a new one for sync tests
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
return AsyncIOScheduler(loop=loop)
|
|
|
|
|
|
@pytest.fixture
|
|
def stream_router(scheduler):
|
|
"""Fixture for ReactiveStreamRouter."""
|
|
return ReactiveStreamRouter(scheduler)
|
|
|
|
|
|
@pytest.fixture
|
|
def bridge(stream_router):
|
|
"""Fixture for RxPyLangGraphBridge."""
|
|
return RxPyLangGraphBridge(stream_router)
|
|
|
|
|
|
@pytest.fixture
|
|
def simple_graph_config():
|
|
"""Fixture for a simple graph configuration."""
|
|
return {
|
|
"name": "test_graph",
|
|
"entry_point": "start",
|
|
"checkpointing": False,
|
|
"enable_time_travel": False,
|
|
"parallel_execution": True,
|
|
"nodes": {
|
|
"start": {
|
|
"type": "start"
|
|
},
|
|
"process": {
|
|
"type": "function",
|
|
"agent": "test_agent"
|
|
},
|
|
"end": {
|
|
"type": "end"
|
|
}
|
|
},
|
|
"edges": [
|
|
{"source": "start", "target": "process"},
|
|
{"source": "process", "target": "end"}
|
|
]
|
|
}
|
|
|
|
|
|
class TestRxPyLangGraphBridgeInit:
|
|
"""Test suite for bridge initialization."""
|
|
|
|
def test_bridge_initialization(self, stream_router):
|
|
"""Test bridge initializes correctly."""
|
|
bridge = RxPyLangGraphBridge(stream_router)
|
|
|
|
assert bridge.stream_router is stream_router
|
|
assert bridge.logger is not None
|
|
assert bridge.graphs == {}
|
|
|
|
def test_bridge_registers_operators(self, stream_router):
|
|
"""Test that bridge registers LangGraph operators."""
|
|
bridge = RxPyLangGraphBridge(stream_router)
|
|
|
|
# Check operator factory methods exist
|
|
assert hasattr(bridge, '_operator_graph_execute')
|
|
assert hasattr(bridge, '_operator_state_update')
|
|
assert hasattr(bridge, '_operator_state_checkpoint')
|
|
assert hasattr(bridge, '_operator_langgraph_node')
|
|
assert hasattr(bridge, '_operator_conditional_route')
|
|
|
|
|
|
class TestGraphCreation:
|
|
"""Test suite for graph creation."""
|
|
|
|
def test_create_graph_from_config_basic(self, bridge, simple_graph_config):
|
|
"""Test creating a basic graph from configuration."""
|
|
graph = bridge.create_graph_from_config(simple_graph_config)
|
|
|
|
assert isinstance(graph, LangGraph)
|
|
assert graph.name == "test_graph"
|
|
assert "test_graph" in bridge.graphs
|
|
|
|
def test_create_graph_with_custom_config(self, bridge):
|
|
"""Test creating graph with custom configuration."""
|
|
config = {
|
|
"name": "custom_graph",
|
|
"entry_point": "custom_start",
|
|
"checkpointing": True,
|
|
"enable_time_travel": True,
|
|
"parallel_execution": False,
|
|
"nodes": {
|
|
"custom_start": {"type": "start"}
|
|
},
|
|
"edges": []
|
|
}
|
|
|
|
graph = bridge.create_graph_from_config(config)
|
|
|
|
assert graph.name == "custom_graph"
|
|
assert graph.config.entry_point == "custom_start"
|
|
assert graph.config.checkpointing is True
|
|
|
|
def test_create_graph_with_nodes(self, bridge):
|
|
"""Test creating graph with various node types."""
|
|
config = {
|
|
"name": "node_graph",
|
|
"nodes": {
|
|
"func_node": {
|
|
"type": "function",
|
|
"agent": "agent1",
|
|
"tools": ["tool1", "tool2"],
|
|
"timeout": 30,
|
|
"parallel": True,
|
|
"metadata": {"key": "value"}
|
|
},
|
|
"cond_node": {
|
|
"type": "conditional",
|
|
"condition": "some_condition"
|
|
}
|
|
},
|
|
"edges": []
|
|
}
|
|
|
|
graph = bridge.create_graph_from_config(config)
|
|
|
|
assert "func_node" in graph.config.nodes
|
|
assert "cond_node" in graph.config.nodes
|
|
|
|
def test_create_graph_with_edges(self, bridge):
|
|
"""Test creating graph with edges."""
|
|
config = {
|
|
"name": "edge_graph",
|
|
"nodes": {
|
|
"node1": {"type": "function"},
|
|
"node2": {"type": "function"}
|
|
},
|
|
"edges": [
|
|
{
|
|
"source": "node1",
|
|
"target": "node2",
|
|
"condition": "some_condition",
|
|
"metadata": {"weight": 1.0}
|
|
}
|
|
]
|
|
}
|
|
|
|
graph = bridge.create_graph_from_config(config)
|
|
|
|
assert len(graph.config.edges) == 1
|
|
assert graph.config.edges[0].source == "node1"
|
|
assert graph.config.edges[0].target == "node2"
|
|
|
|
|
|
class TestStreamCreation:
|
|
"""Test suite for stream creation."""
|
|
|
|
def test_create_graph_stream(self, bridge, simple_graph_config):
|
|
"""Test creating a stream for graph execution."""
|
|
bridge.create_graph_from_config(simple_graph_config)
|
|
stream_config = bridge.create_graph_stream("test_graph")
|
|
|
|
assert isinstance(stream_config, StreamConfig)
|
|
assert stream_config.name == "graph_test_graph"
|
|
assert stream_config.type == StreamType.COLD
|
|
assert len(stream_config.operators) == 1
|
|
assert stream_config.operators[0]["type"] == "graph_execute"
|
|
|
|
def test_create_graph_stream_invalid_name(self, bridge):
|
|
"""Test creating stream for non-existent graph raises error."""
|
|
with pytest.raises(ValueError) as exc_info:
|
|
bridge.create_graph_stream("nonexistent_graph")
|
|
|
|
assert "Graph 'nonexistent_graph' not found" in str(exc_info.value)
|
|
|
|
|
|
class TestOperatorCreation:
|
|
"""Test suite for operator creation."""
|
|
|
|
def test_create_graph_executor_operator(self, bridge, simple_graph_config):
|
|
"""Test creating graph executor operator."""
|
|
bridge.create_graph_from_config(simple_graph_config)
|
|
|
|
params = {"graph": "test_graph"}
|
|
operator = bridge._create_graph_executor(params)
|
|
|
|
assert operator is not None
|
|
assert callable(operator)
|
|
|
|
def test_create_graph_executor_invalid_graph(self, bridge):
|
|
"""Test graph executor with invalid graph name."""
|
|
params = {"graph": "invalid_graph"}
|
|
|
|
with pytest.raises(ValueError) as exc_info:
|
|
bridge._create_graph_executor(params)
|
|
|
|
assert "Invalid graph name" in str(exc_info.value)
|
|
|
|
def test_create_graph_executor_no_graph_name(self, bridge):
|
|
"""Test graph executor without graph name."""
|
|
params = {}
|
|
|
|
with pytest.raises(ValueError) as exc_info:
|
|
bridge._create_graph_executor(params)
|
|
|
|
assert "Invalid graph name" in str(exc_info.value)
|
|
|
|
def test_create_state_updater_operator(self, bridge, simple_graph_config):
|
|
"""Test creating state updater operator."""
|
|
bridge.create_graph_from_config(simple_graph_config)
|
|
|
|
params = {"graph": "test_graph"}
|
|
operator = bridge._create_state_updater(params)
|
|
|
|
assert operator is not None
|
|
assert callable(operator)
|
|
|
|
def test_create_state_updater_invalid_graph(self, bridge):
|
|
"""Test state updater with invalid graph name."""
|
|
params = {"graph": "invalid_graph"}
|
|
|
|
with pytest.raises(ValueError) as exc_info:
|
|
bridge._create_state_updater(params)
|
|
|
|
assert "Invalid graph name" in str(exc_info.value)
|
|
|
|
def test_create_state_checkpointer_operator(self, bridge, simple_graph_config):
|
|
"""Test creating state checkpointer operator."""
|
|
bridge.create_graph_from_config(simple_graph_config)
|
|
|
|
params = {"graph": "test_graph"}
|
|
operator = bridge._create_state_checkpointer(params)
|
|
|
|
assert operator is not None
|
|
assert callable(operator)
|
|
|
|
def test_create_state_checkpointer_invalid_graph(self, bridge):
|
|
"""Test state checkpointer with invalid graph name."""
|
|
params = {"graph": "invalid_graph"}
|
|
|
|
with pytest.raises(ValueError) as exc_info:
|
|
bridge._create_state_checkpointer(params)
|
|
|
|
assert "Invalid graph name" in str(exc_info.value)
|
|
|
|
def test_create_node_operator(self, bridge, simple_graph_config):
|
|
"""Test creating node operator."""
|
|
bridge.create_graph_from_config(simple_graph_config)
|
|
|
|
params = {"graph": "test_graph", "node": "process"}
|
|
operator = bridge._create_node_operator(params)
|
|
|
|
assert operator is not None
|
|
assert callable(operator)
|
|
|
|
def test_create_node_operator_invalid_graph(self, bridge):
|
|
"""Test node operator with invalid graph name."""
|
|
params = {"graph": "invalid_graph", "node": "test_node"}
|
|
|
|
with pytest.raises(ValueError) as exc_info:
|
|
bridge._create_node_operator(params)
|
|
|
|
assert "Invalid graph name" in str(exc_info.value)
|
|
|
|
def test_create_node_operator_invalid_node(self, bridge, simple_graph_config):
|
|
"""Test node operator with invalid node name."""
|
|
bridge.create_graph_from_config(simple_graph_config)
|
|
|
|
params = {"graph": "test_graph", "node": "invalid_node"}
|
|
|
|
with pytest.raises(ValueError) as exc_info:
|
|
bridge._create_node_operator(params)
|
|
|
|
assert "Invalid node name" in str(exc_info.value)
|
|
|
|
def test_create_conditional_router(self, bridge):
|
|
"""Test creating conditional router operator."""
|
|
params = {
|
|
"routes": {
|
|
"route1": {"type": "always"},
|
|
"route2": {"type": "content_contains", "text": "test"}
|
|
},
|
|
"default": "default_route"
|
|
}
|
|
|
|
operator = bridge._create_conditional_router(params)
|
|
|
|
assert operator is not None
|
|
assert callable(operator)
|
|
|
|
|
|
class TestRouteConditionEvaluation:
|
|
"""Test suite for route condition evaluation."""
|
|
|
|
def test_evaluate_always_condition(self, bridge):
|
|
"""Test evaluating 'always' condition."""
|
|
msg = StreamMessage(content="test", metadata={})
|
|
condition = {"type": "always"}
|
|
|
|
result = bridge._evaluate_route_condition(msg, condition)
|
|
|
|
assert result is True
|
|
|
|
def test_evaluate_content_type_condition(self, bridge):
|
|
"""Test evaluating 'content_type' condition."""
|
|
msg = StreamMessage(content="test string", metadata={})
|
|
condition = {"type": "content_type", "value": "str"}
|
|
|
|
result = bridge._evaluate_route_condition(msg, condition)
|
|
|
|
assert result is True
|
|
|
|
def test_evaluate_content_type_condition_mismatch(self, bridge):
|
|
"""Test evaluating 'content_type' condition with mismatch."""
|
|
msg = StreamMessage(content="test", metadata={})
|
|
condition = {"type": "content_type", "value": "int"}
|
|
|
|
result = bridge._evaluate_route_condition(msg, condition)
|
|
|
|
assert result is False
|
|
|
|
def test_evaluate_metadata_has_condition(self, bridge):
|
|
"""Test evaluating 'metadata_has' condition."""
|
|
msg = StreamMessage(content="test", metadata={"key": "value"})
|
|
condition = {"type": "metadata_has", "key": "key"}
|
|
|
|
result = bridge._evaluate_route_condition(msg, condition)
|
|
|
|
assert result is True
|
|
|
|
def test_evaluate_metadata_has_condition_missing(self, bridge):
|
|
"""Test evaluating 'metadata_has' condition with missing key."""
|
|
msg = StreamMessage(content="test", metadata={})
|
|
condition = {"type": "metadata_has", "key": "missing_key"}
|
|
|
|
result = bridge._evaluate_route_condition(msg, condition)
|
|
|
|
assert result is False
|
|
|
|
def test_evaluate_content_contains_condition(self, bridge):
|
|
"""Test evaluating 'content_contains' condition."""
|
|
msg = StreamMessage(content="hello world", metadata={})
|
|
condition = {"type": "content_contains", "text": "world"}
|
|
|
|
result = bridge._evaluate_route_condition(msg, condition)
|
|
|
|
assert result is True
|
|
|
|
def test_evaluate_content_contains_condition_not_found(self, bridge):
|
|
"""Test evaluating 'content_contains' condition when text not found."""
|
|
msg = StreamMessage(content="hello world", metadata={})
|
|
condition = {"type": "content_contains", "text": "missing"}
|
|
|
|
result = bridge._evaluate_route_condition(msg, condition)
|
|
|
|
assert result is False
|
|
|
|
def test_evaluate_content_not_contains_condition(self, bridge):
|
|
"""Test evaluating 'content_not_contains' condition."""
|
|
msg = StreamMessage(content="hello world", metadata={})
|
|
condition = {"type": "content_not_contains", "text": "missing"}
|
|
|
|
result = bridge._evaluate_route_condition(msg, condition)
|
|
|
|
assert result is True
|
|
|
|
def test_evaluate_content_not_contains_condition_found(self, bridge):
|
|
"""Test evaluating 'content_not_contains' condition when text is found."""
|
|
msg = StreamMessage(content="hello world", metadata={})
|
|
condition = {"type": "content_not_contains", "text": "world"}
|
|
|
|
result = bridge._evaluate_route_condition(msg, condition)
|
|
|
|
assert result is False
|
|
|
|
def test_evaluate_unknown_condition_type(self, bridge):
|
|
"""Test evaluating unknown condition type returns False."""
|
|
msg = StreamMessage(content="test", metadata={})
|
|
condition = {"type": "unknown_type"}
|
|
|
|
result = bridge._evaluate_route_condition(msg, condition)
|
|
|
|
assert result is False
|
|
|
|
|
|
class TestStreamGraphConnection:
|
|
"""Test suite for connecting streams and graphs."""
|
|
|
|
def test_connect_stream_to_graph(self, bridge, stream_router, simple_graph_config):
|
|
"""Test connecting stream to graph."""
|
|
# Create graph
|
|
bridge.create_graph_from_config(simple_graph_config)
|
|
|
|
# Create a test stream
|
|
test_stream = Subject()
|
|
stream_router.streams["test_stream"] = test_stream
|
|
stream_router.observables["test_stream"] = test_stream
|
|
|
|
# Connect
|
|
bridge.connect_stream_to_graph("test_stream", "test_graph")
|
|
|
|
# Should not raise error
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_connect_stream_to_graph_on_message(self, bridge, stream_router, simple_graph_config):
|
|
"""Test connect_stream_to_graph on_message callback triggers graph execution."""
|
|
# Create graph
|
|
graph = bridge.create_graph_from_config(simple_graph_config)
|
|
|
|
# Create a test stream
|
|
test_stream = Subject()
|
|
stream_router.streams["test_stream"] = test_stream
|
|
stream_router.observables["test_stream"] = test_stream
|
|
|
|
# Mock the graph's execute method to track calls
|
|
graph.execute = AsyncMock()
|
|
|
|
# Connect stream to graph
|
|
bridge.connect_stream_to_graph("test_stream", "test_graph")
|
|
|
|
# Send a message through the stream
|
|
test_message = StreamMessage(content="test input", metadata={})
|
|
test_stream.on_next(test_message)
|
|
|
|
# Give time for async execution
|
|
await asyncio.sleep(0.1)
|
|
|
|
# Verify graph.execute was called with the message content
|
|
graph.execute.assert_called()
|
|
assert graph.execute.call_count == 1
|
|
|
|
def test_connect_stream_to_graph_invalid_stream(self, bridge, simple_graph_config):
|
|
"""Test connecting non-existent stream raises error."""
|
|
bridge.create_graph_from_config(simple_graph_config)
|
|
|
|
with pytest.raises(ValueError) as exc_info:
|
|
bridge.connect_stream_to_graph("nonexistent_stream", "test_graph")
|
|
|
|
assert "Stream 'nonexistent_stream' not found" in str(exc_info.value)
|
|
|
|
def test_connect_stream_to_graph_invalid_graph(self, bridge, stream_router):
|
|
"""Test connecting stream to non-existent graph raises error."""
|
|
test_stream = Subject()
|
|
stream_router.streams["test_stream"] = test_stream
|
|
stream_router.observables["test_stream"] = test_stream
|
|
|
|
with pytest.raises(ValueError) as exc_info:
|
|
bridge.connect_stream_to_graph("test_stream", "nonexistent_graph")
|
|
|
|
assert "Graph 'nonexistent_graph' not found" in str(exc_info.value)
|
|
|
|
def test_connect_graph_to_stream(self, bridge, simple_graph_config):
|
|
"""Test connecting graph to stream."""
|
|
bridge.create_graph_from_config(simple_graph_config)
|
|
|
|
# Connect graph to new stream
|
|
bridge.connect_graph_to_stream("test_graph", "output_stream")
|
|
|
|
# Stream should be created
|
|
assert "output_stream" in bridge.stream_router.streams
|
|
|
|
def test_connect_graph_to_stream_invalid_graph(self, bridge):
|
|
"""Test connecting non-existent graph raises error."""
|
|
with pytest.raises(ValueError) as exc_info:
|
|
bridge.connect_graph_to_stream("nonexistent_graph", "output_stream")
|
|
|
|
assert "Graph 'nonexistent_graph' not found" in str(exc_info.value)
|
|
|
|
def test_connect_graph_to_existing_stream(self, bridge, stream_router, simple_graph_config):
|
|
"""Test connecting graph to existing stream."""
|
|
bridge.create_graph_from_config(simple_graph_config)
|
|
|
|
# Create existing stream
|
|
existing_stream = Subject()
|
|
stream_router.streams["existing_stream"] = existing_stream
|
|
|
|
# Connect graph to existing stream
|
|
bridge.connect_graph_to_stream("test_graph", "existing_stream")
|
|
|
|
# Should use existing stream
|
|
assert stream_router.streams["existing_stream"] is existing_stream
|
|
|
|
|
|
class TestHybridPipeline:
|
|
"""Test suite for hybrid pipeline creation."""
|
|
|
|
def test_create_hybrid_pipeline_empty(self, bridge):
|
|
"""Test creating empty hybrid pipeline."""
|
|
config = {"stages": []}
|
|
|
|
bridge.create_hybrid_pipeline(config)
|
|
|
|
# Should complete without error
|
|
|
|
def test_create_hybrid_pipeline_with_stream_stage(self, bridge):
|
|
"""Test creating hybrid pipeline with stream stage."""
|
|
config = {
|
|
"stages": [
|
|
{
|
|
"type": "stream",
|
|
"name": "stream1",
|
|
"stream_type": "cold",
|
|
"operators": [],
|
|
"publications": []
|
|
}
|
|
]
|
|
}
|
|
|
|
bridge.create_hybrid_pipeline(config)
|
|
|
|
# Stream should be created
|
|
assert "stream1" in bridge.stream_router.streams
|
|
|
|
def test_create_hybrid_pipeline_with_graph_stage(self, bridge):
|
|
"""Test creating hybrid pipeline with graph stage."""
|
|
config = {
|
|
"stages": [
|
|
{
|
|
"type": "graph",
|
|
"config": {
|
|
"name": "graph1",
|
|
"nodes": {},
|
|
"edges": []
|
|
}
|
|
}
|
|
]
|
|
}
|
|
|
|
bridge.create_hybrid_pipeline(config)
|
|
|
|
# Graph should be created
|
|
assert "graph1" in bridge.graphs
|
|
|
|
def test_create_hybrid_pipeline_with_connections(self, bridge, stream_router):
|
|
"""Test creating hybrid pipeline with stage connections."""
|
|
# Create input stream first
|
|
input_stream = Subject()
|
|
stream_router.streams["input_stream"] = input_stream
|
|
stream_router.observables["input_stream"] = input_stream
|
|
|
|
config = {
|
|
"stages": [
|
|
{
|
|
"type": "stream",
|
|
"name": "stream1",
|
|
"stream_type": "cold",
|
|
"operators": [],
|
|
"publications": []
|
|
},
|
|
{
|
|
"type": "graph",
|
|
"config": {
|
|
"name": "graph1",
|
|
"nodes": {},
|
|
"edges": []
|
|
},
|
|
"input_from": "input_stream",
|
|
"output_to": "output_stream"
|
|
}
|
|
]
|
|
}
|
|
|
|
bridge.create_hybrid_pipeline(config)
|
|
|
|
# Graph and streams should be created
|
|
assert "graph1" in bridge.graphs
|
|
assert "output_stream" in bridge.stream_router.streams
|
|
|
|
|
|
class TestUtilityMethods:
|
|
"""Test suite for utility methods."""
|
|
|
|
def test_run_async_safely(self, bridge):
|
|
"""Test running async coroutine safely."""
|
|
async def test_coro():
|
|
return "result"
|
|
|
|
coro = test_coro()
|
|
result = bridge._run_async_safely(coro)
|
|
|
|
# Should return the coroutine itself
|
|
assert asyncio.iscoroutine(result)
|
|
coro.close() # Clean up
|
|
|
|
def test_get_graph(self, bridge, simple_graph_config):
|
|
"""Test getting graph by name."""
|
|
graph = bridge.create_graph_from_config(simple_graph_config)
|
|
|
|
retrieved = bridge.get_graph("test_graph")
|
|
|
|
assert retrieved is graph
|
|
|
|
def test_get_graph_nonexistent(self, bridge):
|
|
"""Test getting non-existent graph returns None."""
|
|
result = bridge.get_graph("nonexistent")
|
|
|
|
assert result is None
|
|
|
|
def test_list_graphs_empty(self, bridge):
|
|
"""Test listing graphs when empty."""
|
|
graphs = bridge.list_graphs()
|
|
|
|
assert graphs == []
|
|
|
|
def test_list_graphs(self, bridge, simple_graph_config):
|
|
"""Test listing graphs."""
|
|
bridge.create_graph_from_config(simple_graph_config)
|
|
bridge.create_graph_from_config({**simple_graph_config, "name": "graph2"})
|
|
|
|
graphs = bridge.list_graphs()
|
|
|
|
assert "test_graph" in graphs
|
|
assert "graph2" in graphs
|
|
assert len(graphs) == 2
|
|
|
|
|
|
class TestAsyncExecution:
|
|
"""Test suite for async execution paths."""
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("has_messages,expected_content_type", [
|
|
(True, str), # With messages
|
|
(False, (dict, str)) # Without messages (uses to_dict)
|
|
])
|
|
async def test_graph_executor_execution(self, bridge, simple_graph_config, has_messages, expected_content_type):
|
|
"""Test graph executor execution with and without messages."""
|
|
# Create graph with mocked agent
|
|
mock_agent = Mock()
|
|
mock_agent.process_message = AsyncMock(return_value="processed")
|
|
bridge.stream_router.agents = {"test_agent": mock_agent}
|
|
|
|
graph = bridge.create_graph_from_config(simple_graph_config)
|
|
|
|
# Mock graph execute method
|
|
async def mock_execute(input_data):
|
|
state = GraphState()
|
|
if has_messages:
|
|
state.messages = [{"role": "assistant", "content": "result"}]
|
|
else:
|
|
state.messages = []
|
|
state.data = {"key": "value", "result": "test_data"}
|
|
return state
|
|
|
|
graph.execute = AsyncMock(side_effect=mock_execute)
|
|
graph.get_execution_history = Mock(return_value=["start", "process", "end"])
|
|
|
|
# Create operator
|
|
params = {"graph": "test_graph"}
|
|
operator = bridge._create_graph_executor(params)
|
|
|
|
# Test with string message
|
|
msg = StreamMessage(content="test input", metadata={"key": "value"})
|
|
|
|
# Execute through RxPy
|
|
result_messages = []
|
|
|
|
def on_next(msg):
|
|
result_messages.append(msg)
|
|
|
|
# Create source and apply operator
|
|
source = rx.of(msg)
|
|
_ = source.pipe(operator).subscribe(on_next=on_next)
|
|
|
|
# Wait a bit for async execution
|
|
await asyncio.sleep(0.1)
|
|
|
|
# Check that execution happened
|
|
assert graph.execute.call_count == 1
|
|
assert len(result_messages) == 1
|
|
if result_messages:
|
|
assert isinstance(result_messages[0].content, expected_content_type)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_state_updater_execution(self, bridge, simple_graph_config):
|
|
"""Test execution of state updater operator."""
|
|
graph = bridge.create_graph_from_config(simple_graph_config)
|
|
|
|
# Mock state manager
|
|
graph.state_manager.update_state = Mock()
|
|
|
|
params = {"graph": "test_graph"}
|
|
operator = bridge._create_state_updater(params)
|
|
|
|
# Test with dict content
|
|
msg = StreamMessage(content={"data": "value"}, metadata={"orig": "meta"})
|
|
|
|
result_messages = []
|
|
|
|
source = rx.of(msg)
|
|
_ = source.pipe(operator).subscribe(on_next= result_messages.append)
|
|
|
|
assert len(result_messages) == 1
|
|
assert result_messages[0].metadata["state_updated"] is True
|
|
assert graph.state_manager.update_state.called
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_state_updater_with_string_content(self, bridge, simple_graph_config):
|
|
"""Test state updater with string content."""
|
|
graph = bridge.create_graph_from_config(simple_graph_config)
|
|
graph.state_manager.update_state = Mock()
|
|
|
|
params = {"graph": "test_graph"}
|
|
operator = bridge._create_state_updater(params)
|
|
|
|
# Test with string content (should be wrapped)
|
|
msg = StreamMessage(content="string data", metadata={})
|
|
|
|
result_messages = []
|
|
|
|
source = rx.of(msg)
|
|
_ = source.pipe(operator).subscribe(on_next= result_messages.append)
|
|
|
|
assert len(result_messages) == 1
|
|
# Should have called update_state with wrapped content
|
|
graph.state_manager.update_state.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_state_checkpointer_execution(self, bridge, simple_graph_config):
|
|
"""Test execution of state checkpointer operator."""
|
|
graph = bridge.create_graph_from_config(simple_graph_config)
|
|
|
|
# Mock the _save_checkpoint method
|
|
graph.state_manager._save_checkpoint = Mock()
|
|
|
|
params = {"graph": "test_graph"}
|
|
operator = bridge._create_state_checkpointer(params)
|
|
|
|
msg = StreamMessage(content="test", metadata={})
|
|
|
|
result_messages = []
|
|
|
|
source = rx.of(msg)
|
|
_ = source.pipe(operator).subscribe(on_next= result_messages.append)
|
|
|
|
assert len(result_messages) == 1
|
|
assert result_messages[0].metadata["checkpointed"] is True
|
|
assert graph.state_manager._save_checkpoint.called
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_node_operator_execution(self, bridge, simple_graph_config):
|
|
"""Test execution of node operator."""
|
|
# Create graph with mocked agent
|
|
mock_agent = Mock()
|
|
mock_agent.process_message = AsyncMock(return_value="node result")
|
|
bridge.stream_router.agents = {"test_agent": mock_agent}
|
|
|
|
graph = bridge.create_graph_from_config(simple_graph_config)
|
|
|
|
# Mock node execution
|
|
process_node = graph.nodes["process"]
|
|
async def mock_node_execute(state):
|
|
return {"messages": [{"role": "assistant", "content": "node output"}]}
|
|
|
|
process_node.execute = AsyncMock(side_effect=mock_node_execute)
|
|
graph.state_manager.update_state = Mock()
|
|
|
|
params = {"graph": "test_graph", "node": "process"}
|
|
operator = bridge._create_node_operator(params)
|
|
|
|
msg = StreamMessage(content="test input", metadata={})
|
|
|
|
# Execute
|
|
result_messages = []
|
|
|
|
source = rx.of(msg)
|
|
_ = source.pipe(operator).subscribe(on_next= result_messages.append)
|
|
|
|
# Wait for async execution
|
|
await asyncio.sleep(0.1)
|
|
|
|
# Check execution happened
|
|
assert process_node.execute.call_count == 1
|
|
assert len(result_messages) == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_node_operator_with_string_message(self, bridge, simple_graph_config):
|
|
"""Test node operator with string message content."""
|
|
mock_agent = Mock()
|
|
mock_agent.process_message = AsyncMock(return_value="result")
|
|
bridge.stream_router.agents = {"test_agent": mock_agent}
|
|
|
|
graph = bridge.create_graph_from_config(simple_graph_config)
|
|
process_node = graph.nodes["process"]
|
|
|
|
# Mock to return updates without messages
|
|
async def mock_execute(state):
|
|
return {"data": "some_data"}
|
|
|
|
process_node.execute = AsyncMock(side_effect=mock_execute)
|
|
graph.state_manager.update_state = Mock()
|
|
|
|
params = {"graph": "test_graph", "node": "process"}
|
|
operator = bridge._create_node_operator(params)
|
|
|
|
msg = StreamMessage(content="test", metadata={"key": "val"})
|
|
|
|
result_messages = []
|
|
|
|
source = rx.of(msg)
|
|
_ = source.pipe(operator).subscribe(on_next= result_messages.append)
|
|
|
|
await asyncio.sleep(0.1)
|
|
|
|
# Should have processed
|
|
assert process_node.execute.call_count == 1
|
|
assert len(result_messages) == 1
|
|
|
|
def test_conditional_router_execution(self, bridge):
|
|
"""Test execution of conditional router."""
|
|
params = {
|
|
"routes": {
|
|
"route_a": {"type": "content_contains", "text": "hello"},
|
|
"route_b": {"type": "content_contains", "text": "goodbye"},
|
|
123: {"type": "always"} # Non-string route name
|
|
},
|
|
"default": "default_route"
|
|
}
|
|
|
|
router_operator = bridge._create_conditional_router(params)
|
|
|
|
# Test routing
|
|
messages = [
|
|
StreamMessage(content="hello world", metadata={}),
|
|
StreamMessage(content="goodbye world", metadata={}),
|
|
StreamMessage(content="other content", metadata={}),
|
|
]
|
|
|
|
routed_messages = []
|
|
|
|
source = rx.from_(messages)
|
|
_ = source.pipe(router_operator).subscribe(on_next= routed_messages.append)
|
|
|
|
# Should have routed all messages
|
|
assert len(routed_messages) > 0
|
|
|
|
def test_conditional_router_with_non_string_default(self, bridge):
|
|
"""Test conditional router with non-string default route."""
|
|
params = {
|
|
"routes": {},
|
|
"default": 456 # Non-string default
|
|
}
|
|
|
|
router_operator = bridge._create_conditional_router(params)
|
|
|
|
# The operator should handle this gracefully
|
|
assert router_operator is not None
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"])
|
|
|