"""Pytest fixtures and configuration.""" import sys from pathlib import Path import pytest import torch sys.path.insert(0, str(Path(__file__).parent.parent / "src")) import rustworkx as rx @pytest.fixture def empty_digraph(): """Empty directed graph.""" return rx.PyDiGraph() @pytest.fixture def simple_digraph(): """Simple directed graph with 3 nodes.""" g = rx.PyDiGraph() g.add_node({"id": "a", "role": "agent"}) g.add_node({"id": "b", "role": "agent"}) g.add_node({"id": "c", "role": "agent"}) g.add_edge(0, 1, {"weight": 0.8}) g.add_edge(1, 2, {"weight": 0.9}) return g @pytest.fixture def cyclic_digraph(): """Graph with a cycle.""" g = rx.PyDiGraph() g.add_node({"id": "a"}) g.add_node({"id": "b"}) g.add_node({"id": "c"}) g.add_edge(0, 1, {"weight": 1.0}) g.add_edge(1, 2, {"weight": 1.0}) g.add_edge(2, 0, {"weight": 1.0}) # Cycle return g @pytest.fixture def parallel_digraph(): """Graph with parallel branches.""" g = rx.PyDiGraph() # Task -> A, B (parallel) -> C g.add_node({"id": "task", "role": "task"}) g.add_node({"id": "a", "role": "agent"}) g.add_node({"id": "b", "role": "agent"}) g.add_node({"id": "c", "role": "agent"}) g.add_edge(0, 1, {"weight": 0.5}) g.add_edge(0, 2, {"weight": 0.5}) g.add_edge(1, 3, {"weight": 0.8}) g.add_edge(2, 3, {"weight": 0.8}) return g @pytest.fixture def sample_agent_configs(): """Sample agent configurations.""" return [ { "agent_id": "coordinator", "role_name": "Coordinator", "persona": "Manages workflow", "state": {}, }, { "agent_id": "researcher", "role_name": "Researcher", "persona": "Finds information", "state": {}, }, { "agent_id": "writer", "role_name": "Writer", "persona": "Creates content", "state": {}, }, ] @pytest.fixture def sample_connections(): """Sample agent connections.""" return { "coordinator": ["researcher", "writer"], "researcher": ["writer"], "writer": [], } @pytest.fixture def mock_agent_callable(): """Mock callable for an agent.""" async def agent_fn(query: str, context: dict) -> str: del context # Unused in mock return f"Response to: {query}" return agent_fn @pytest.fixture def mock_failing_agent(): """Mock agent that always fails.""" async def failing_agent(query: str, context: dict) -> str: del query, context # Unused in mock msg = "Agent failure" raise RuntimeError(msg) return failing_agent @pytest.fixture def mock_slow_agent(): """Mock slow agent.""" import asyncio async def slow_agent(query: str, context: dict) -> str: del query, context # Unused in mock await asyncio.sleep(5.0) # Intentionally slow return "Slow response" return slow_agent def create_role_graph(node_ids, connections): """Helper for creating a RoleGraph.""" from core.agent import AgentProfile from core.graph import RoleGraph g = rx.PyDiGraph() id_to_idx = {} agents = [] for nid in node_ids: idx = g.add_node({"id": nid}) id_to_idx[nid] = idx agent = AgentProfile(agent_id=nid, display_name=f"Agent {nid.upper()}") agents.append(agent) for src, targets in connections.items(): for tgt in targets: if src in id_to_idx and tgt in id_to_idx: g.add_edge(id_to_idx[src], id_to_idx[tgt], {"weight": 1.0}) n = len(node_ids) a_com = torch.zeros((n, n), dtype=torch.float32) for src, targets in connections.items(): i = node_ids.index(src) for tgt in targets: if tgt in node_ids: j = node_ids.index(tgt) a_com[i, j] = 1.0 graph = RoleGraph( node_ids=node_ids, role_connections=connections, graph=g, A_com=a_com, ) graph.agents = agents return graph