"""Steps for remaining uncovered lines/branches in bridge.py. Targets: - Line 35: ``cancellation_reasons`` WeakKeyDictionary initialisation - Line 68: ``task.cancel()`` inside ``cleanup_tasks_async`` - Lines 201-204: string-content branch in the inner ``execute_graph`` coroutine - Branches 39->38 (__del__ normal exit), 67->68 (not-done cancel in async cleanup), 182->184 (valid graph in create_graph_stream), 260->258 (valid graph in checkpointer) """ from __future__ import annotations import asyncio from typing import Any from unittest.mock import AsyncMock, MagicMock from behave import given, then, when from behave.runner import Context from cleveragents.langgraph.bridge import RxPyLangGraphBridge from cleveragents.langgraph.state import GraphState from cleveragents.reactive.stream_router import ( ReactiveStreamRouter, StreamMessage, ) # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _build_bridge() -> tuple[RxPyLangGraphBridge, MagicMock]: """Return a bridge backed by a mocked stream-router.""" scheduler = MagicMock() router = ReactiveStreamRouter(scheduler=scheduler) router.agents = {"agent": MagicMock()} bridge = RxPyLangGraphBridge(router) return bridge, scheduler def _build_mock_graph( name: str, *, messages: list[dict[str, Any]] | None = None, state_dict: dict[str, Any] | None = None, ) -> MagicMock: """Build a mock LangGraph whose ``execute`` is an AsyncMock.""" graph = MagicMock() graph.name = name gs = MagicMock(spec=GraphState) gs.messages = messages if messages is not None else [] gs.to_dict = MagicMock(return_value=state_dict or {"messages": gs.messages}) graph.execute = AsyncMock(return_value=gs) graph.get_execution_history = MagicMock(return_value=["s1"]) graph.state_manager = MagicMock() graph.state_manager.checkpoint_dir = None graph.nodes = {} return graph # =================================================================== # Scenario: Cancellation reason is recorded when cancelling a tracked task # =================================================================== @given("a bridge instance prepared for cancellation reason tracking") def step_bridge_for_cancel_reason(context: Context) -> None: bridge, _ = _build_bridge() # Create an event loop and a dummy task loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) async def _forever() -> None: await asyncio.sleep(9999) task = loop.create_task(_forever()) bridge._active_tasks.add(task) context.bridge = bridge context.dummy_task = task context.cancel_loop = loop context.caught_error = None def _cleanup() -> None: # Cancel all pending tasks and close the loop to avoid __del__ errors for t in list(bridge._active_tasks): t.cancel() bridge._active_tasks.clear() if not loop.is_closed(): loop.run_until_complete(asyncio.sleep(0)) loop.close() asyncio.set_event_loop(asyncio.new_event_loop()) context._cleanup_handlers.append(_cleanup) @when('I cancel a tracked task providing reason "{reason}"') def step_cancel_task_with_given_reason(context: Context, reason: str) -> None: loop = context.cancel_loop async def _do_cancel() -> None: await context.bridge.cancel_task_with_reason(context.dummy_task, reason) loop.run_until_complete(_do_cancel()) context.cancel_reason_used = reason @then("the bridge cancellation_reasons mapping should contain the task") def step_assert_cancellation_reasons_has_task(context: Context) -> None: assert context.dummy_task in context.bridge.cancellation_reasons, ( "Expected task to be present in cancellation_reasons" ) @then('the stored reason should equal "{expected}"') def step_assert_stored_reason(context: Context, expected: str) -> None: actual = context.bridge.cancellation_reasons.get(context.dummy_task) assert actual == expected, f"Expected reason '{expected}', got '{actual}'" # Cleanup context.cancel_loop.close() asyncio.set_event_loop(asyncio.new_event_loop()) # =================================================================== # Scenario: Cancelling a None task raises ValueError # =================================================================== @when('I try cancelling a None task with reason "{reason}"') def step_cancel_none_task(context: Context, reason: str) -> None: loop = context.cancel_loop async def _do_cancel() -> None: await context.bridge.cancel_task_with_reason(None, reason) # type: ignore[arg-type] try: loop.run_until_complete(_do_cancel()) except ValueError as exc: context.caught_error = exc @then('a ValueError mentioning "{fragment}" should be raised') def step_then_valueerror_mentioning(context: Context, fragment: str) -> None: assert context.caught_error is not None, "Expected a ValueError but none was raised" assert isinstance(context.caught_error, ValueError), ( f"Expected ValueError, got {type(context.caught_error).__name__}" ) assert fragment in str(context.caught_error), ( f"Expected '{fragment}' in error message: {context.caught_error}" ) # =================================================================== # Scenario: Cancelling a task with empty reason raises ValueError # =================================================================== @when("I try cancelling a valid task with an empty reason string") def step_cancel_task_empty_reason(context: Context) -> None: loop = context.cancel_loop async def _do_cancel() -> None: await context.bridge.cancel_task_with_reason(context.dummy_task, "") try: loop.run_until_complete(_do_cancel()) except ValueError as exc: context.caught_error = exc # =================================================================== # Scenario: Async cleanup cancels a not-yet-done task (line 68, branch 67->68) # =================================================================== @given("a bridge holding a deliberately stalled coroutine task") def step_bridge_with_stalled_task(context: Context) -> None: bridge, _ = _build_bridge() loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) async def _stall() -> None: await asyncio.sleep(99999) task = loop.create_task(_stall()) bridge._active_tasks.add(task) context.bridge = bridge context.stalled_task = task context.async_cleanup_loop = loop @when("I perform async cleanup with a 200ms timeout on stalled tasks") def step_perform_async_cleanup(context: Context) -> None: loop = context.async_cleanup_loop loop.run_until_complete(context.bridge.cleanup_tasks_async(timeout=0.2)) @then("every stalled task should have been cancelled or finished") def step_assert_stalled_tasks_cancelled(context: Context) -> None: task = context.stalled_task assert task.cancelled() or task.done(), ( f"Expected stalled task to be cancelled/done, state={task._state}" ) @then("the bridge task tracking set should be empty after async cleanup") def step_assert_tracking_set_empty(context: Context) -> None: assert len(context.bridge._active_tasks) == 0, ( f"Expected empty _active_tasks, got {len(context.bridge._active_tasks)}" ) # Cleanup loop if ( hasattr(context, "async_cleanup_loop") and not context.async_cleanup_loop.is_closed() ): context.async_cleanup_loop.close() asyncio.set_event_loop(asyncio.new_event_loop()) # =================================================================== # Scenario: Async cleanup handles already-completed tasks # =================================================================== @given("a bridge holding an already-finished async task") def step_bridge_with_finished_task(context: Context) -> None: bridge, _ = _build_bridge() loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) async def _instant() -> str: return "done" task = loop.create_task(_instant()) # Let the task complete loop.run_until_complete(task) assert task.done(), "Task should be done before adding to bridge" bridge._active_tasks.add(task) context.bridge = bridge context.async_cleanup_loop = loop # =================================================================== # Scenario: Graph executor inner coroutine with string content (lines 201-204) # =================================================================== @given("a bridge with a mock graph wired for direct coroutine invocation") def step_bridge_for_direct_coroutine(context: Context) -> None: bridge, _ = _build_bridge() mock_graph = _build_mock_graph( "direct_exec", messages=[{"content": "coroutine-ok"}], state_dict={"messages": [{"content": "coroutine-ok"}]}, ) bridge.graphs["direct_exec"] = mock_graph context.bridge = bridge context.direct_graph = mock_graph context.exec_loop = asyncio.new_event_loop() asyncio.set_event_loop(context.exec_loop) @when("I invoke the executor coroutine with a plain string message body") def step_invoke_executor_with_string(context: Context) -> None: """Push a string message through the executor and await the task directly.""" import rx # type: ignore bridge = context.bridge executor_op = bridge._create_graph_executor({"graph": "direct_exec"}) loop = context.exec_loop results: list[Any] = [] errors: list[Any] = [] msg = StreamMessage(content="hello from string", metadata={"src": "test"}) rx.just(msg).pipe(executor_op).subscribe( on_next=lambda x: results.append(x), on_error=lambda e: errors.append(e), ) # Drain all pending tasks by awaiting them directly pending = list(bridge._active_tasks) if pending: loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) else: loop.run_until_complete(asyncio.sleep(0.05)) context.coroutine_results = results context.coroutine_errors = errors @then("the mock graph execute should have received a messages list with the string") def step_assert_graph_received_string(context: Context) -> None: assert len(context.coroutine_errors) == 0, ( f"Coroutine raised errors: {context.coroutine_errors}" ) # Verify that graph.execute was called with input_data containing messages call_args = context.direct_graph.execute.call_args assert call_args is not None, "graph.execute was never called" input_data = call_args[0][0] # First positional argument assert "messages" in input_data, ( f"Expected 'messages' key in input_data, got keys: {list(input_data.keys())}" ) assert input_data["messages"][0]["content"] == "hello from string", ( f"Expected string content in messages, got: {input_data['messages']}" ) @then("the executor coroutine should return a StreamMessage with graph metadata") def step_assert_coroutine_returns_stream_message(context: Context) -> None: assert len(context.coroutine_results) > 0, "No results from executor coroutine" result_msg = context.coroutine_results[0] assert isinstance(result_msg, StreamMessage), ( f"Expected StreamMessage, got {type(result_msg).__name__}" ) assert "graph" in result_msg.metadata, ( f"Expected 'graph' in metadata, keys: {list(result_msg.metadata.keys())}" ) # Cleanup if hasattr(context, "exec_loop") and not context.exec_loop.is_closed(): context.exec_loop.close() asyncio.set_event_loop(asyncio.new_event_loop()) # =================================================================== # Scenario: Graph executor inner coroutine with list content (else branch) # =================================================================== @when("I invoke the executor coroutine with a list payload as message body") def step_invoke_executor_with_list(context: Context) -> None: """Push a list (non-str, non-dict) message through the executor.""" import rx # type: ignore bridge = context.bridge executor_op = bridge._create_graph_executor({"graph": "direct_exec"}) loop = context.exec_loop results: list[Any] = [] errors: list[Any] = [] list_payload = ["item_a", "item_b", "item_c"] msg = StreamMessage(content=list_payload, metadata={}) rx.just(msg).pipe(executor_op).subscribe( on_next=lambda x: results.append(x), on_error=lambda e: errors.append(e), ) pending = list(bridge._active_tasks) if pending: loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) else: loop.run_until_complete(asyncio.sleep(0.05)) context.coroutine_results = results context.coroutine_errors = errors @then("the mock graph execute should have received a content key wrapping the list") def step_assert_graph_received_list_wrapped(context: Context) -> None: assert len(context.coroutine_errors) == 0, ( f"Coroutine raised errors: {context.coroutine_errors}" ) call_args = context.direct_graph.execute.call_args assert call_args is not None, "graph.execute was never called" input_data = call_args[0][0] assert "content" in input_data, ( f"Expected 'content' key in input_data for list fallback, got: {list(input_data.keys())}" ) assert input_data["content"] == ["item_a", "item_b", "item_c"], ( f"Expected list content, got: {input_data['content']}" ) # Cleanup if hasattr(context, "exec_loop") and not context.exec_loop.is_closed(): context.exec_loop.close() asyncio.set_event_loop(asyncio.new_event_loop()) # =================================================================== # Scenario: Graph stream config with correct type/publication (branch 182->184) # =================================================================== @given('a bridge with a registered graph called "{name}"') def step_bridge_with_named_registered_graph(context: Context, name: str) -> None: bridge, _ = _build_bridge() mock_graph = _build_mock_graph(name) bridge.graphs[name] = mock_graph context.bridge = bridge @when('I obtain the stream configuration for graph "{name}"') def step_obtain_stream_config(context: Context, name: str) -> None: context.obtained_stream_config = context.bridge.create_graph_stream(name) @then('the stream config name should equal "{expected}"') def step_assert_stream_name_equals(context: Context, expected: str) -> None: actual = context.obtained_stream_config.name assert actual == expected, f"Expected stream name '{expected}', got '{actual}'" @then('the stream config publications should include "{publication}"') def step_assert_publications_include(context: Context, publication: str) -> None: pubs = context.obtained_stream_config.publications assert publication in pubs, f"Expected '{publication}' in publications {pubs}" # =================================================================== # Scenario: State checkpointer operator for valid graph (branch 260->258) # =================================================================== @given('the graph "{name}" has a checkpoint directory configured') def step_configure_checkpoint_dir(context: Context, name: str) -> None: graph = context.bridge.graphs[name] graph.state_manager.checkpoint_dir = "/tmp/test_ckpt" @when('I construct and apply the checkpointer operator for "{name}"') def step_construct_and_apply_checkpointer(context: Context, name: str) -> None: import rx # type: ignore checkpointer_op = context.bridge._create_state_checkpointer({"graph": name}) msg = StreamMessage(content="ckpt-payload", metadata={"check": True}) results: list[Any] = [] rx.just(msg).pipe(checkpointer_op).subscribe( on_next=lambda x: results.append(x), ) context.ckpt_results = results @then("the checkpointer should have invoked _save_checkpoint on the state manager") def step_assert_save_checkpoint_called(context: Context) -> None: assert len(context.ckpt_results) > 0, "Checkpointer produced no output" # The graph "ckpt_valid" should have had _save_checkpoint called graph = context.bridge.graphs.get("ckpt_valid") if graph is not None: graph.state_manager._save_checkpoint.assert_called() # =================================================================== # Scenario: __del__ completes normally (branch 39->38 normal exit) # =================================================================== @given("a fully functional bridge for destructor testing") def step_functional_bridge_for_del(context: Context) -> None: bridge, _ = _build_bridge() context.bridge = bridge context.del_call_error = None @when("__del__ is called on the fully functional bridge") def step_call_del_on_functional_bridge(context: Context) -> None: try: context.bridge.__del__() except Exception as exc: context.del_call_error = exc @then("the __del__ call should complete without any error") def step_assert_del_no_error(context: Context) -> None: assert context.del_call_error is None, ( f"Expected no error from __del__, got {context.del_call_error!r}" )