"""HandoffManager 单元测试""" import asyncio import json import pytest from agentkit.core.protocol import HandoffMessage from agentkit.orchestrator.handoff import HandoffManager # ── HandoffMessage 创建与序列化测试 ───────────────────────────── class TestHandoffMessage: """HandoffMessage 创建与序列化测试""" def test_creation_with_required_fields(self): msg = HandoffMessage( source_agent="agent_a", target_agent="agent_b", task_id="task-001", task_type="analysis", context={"key": "value"}, reason="needs expertise", ) assert msg.source_agent == "agent_a" assert msg.target_agent == "agent_b" assert msg.task_id == "task-001" assert msg.task_type == "analysis" assert msg.context == {"key": "value"} assert msg.reason == "needs expertise" assert msg.created_at is not None def test_to_dict_roundtrip(self): msg = HandoffMessage( source_agent="agent_a", target_agent="agent_b", task_id="task-001", task_type="analysis", context={"data": [1, 2, 3]}, reason="specialization", ) d = msg.to_dict() restored = HandoffMessage.from_dict(d) assert restored.source_agent == msg.source_agent assert restored.target_agent == msg.target_agent assert restored.task_id == msg.task_id assert restored.task_type == msg.task_type assert restored.context == msg.context assert restored.reason == msg.reason def test_to_dict_contains_all_fields(self): msg = HandoffMessage( source_agent="a", target_agent="b", task_id="t1", task_type="search", context={"q": "test"}, reason="handoff", ) d = msg.to_dict() assert "source_agent" in d assert "target_agent" in d assert "task_id" in d assert "task_type" in d assert "context" in d assert "reason" in d assert "created_at" in d def test_from_dict_defaults_context(self): data = { "source_agent": "a", "target_agent": "b", "task_id": "t1", "task_type": "search", "reason": "test", } msg = HandoffMessage.from_dict(data) assert msg.context == {} def test_from_dict_parses_created_at_string(self): data = { "source_agent": "a", "target_agent": "b", "task_id": "t1", "task_type": "search", "context": {}, "reason": "test", "created_at": "2025-01-15T10:30:00+00:00", } msg = HandoffMessage.from_dict(data) assert msg.created_at.year == 2025 assert msg.created_at.month == 1 assert msg.created_at.day == 15 def test_json_serializable(self): msg = HandoffMessage( source_agent="agent_a", target_agent="agent_b", task_id="task-001", task_type="analysis", context={"key": "value"}, reason="needs expertise", ) serialized = json.dumps(msg.to_dict()) deserialized = json.loads(serialized) restored = HandoffMessage.from_dict(deserialized) assert restored.source_agent == msg.source_agent assert restored.target_agent == msg.target_agent assert restored.task_id == msg.task_id # ── HandoffManager 无 Redis（本地模式）测试 ────────────────────── class TestHandoffManagerLocalMode: """HandoffManager 无 Redis（本地模式）测试""" def test_construction_without_redis(self): manager = HandoffManager() assert manager._redis is None assert manager._handlers == {} def test_construction_with_dispatcher(self): manager = HandoffManager(dispatcher="mock_dispatcher") assert manager._dispatcher == "mock_dispatcher" async def test_send_handoff_without_redis_raises(self): manager = HandoffManager() handoff = HandoffMessage( source_agent="a", target_agent="b", task_id="t1", task_type="search", context={}, reason="test", ) with pytest.raises(RuntimeError, match="Redis connection"): await manager.send_handoff(handoff) async def test_listen_for_handoffs_without_redis_returns(self): manager = HandoffManager() # 无 Redis 时应直接返回，不报错 await manager.listen_for_handoffs("agent_a") def test_register_handler(self): manager = HandoffManager() async def handler(msg): pass manager.register_handler("agent_a", handler) assert "agent_a" in manager._handlers assert handler in manager._handlers["agent_a"] def test_register_multiple_handlers_for_same_agent(self): manager = HandoffManager() async def handler1(msg): pass async def handler2(msg): pass manager.register_handler("agent_a", handler1) manager.register_handler("agent_a", handler2) assert len(manager._handlers["agent_a"]) == 2 def test_register_handlers_for_different_agents(self): manager = HandoffManager() async def handler_a(msg): pass async def handler_b(msg): pass manager.register_handler("agent_a", handler_a) manager.register_handler("agent_b", handler_b) assert "agent_a" in manager._handlers assert "agent_b" in manager._handlers assert len(manager._handlers) == 2 # ── HandoffManager _handle_handoff 测试 ───────────────────────── class TestHandoffManagerHandleHandoff: """HandoffManager 内部 _handle_handoff 测试""" async def test_handle_handoff_calls_registered_handlers(self): manager = HandoffManager() received = [] async def handler(msg): received.append(msg) manager.register_handler("agent_b", handler) handoff = HandoffMessage( source_agent="agent_a", target_agent="agent_b", task_id="t1", task_type="search", context={"q": "test"}, reason="delegation", ) await manager._handle_handoff(handoff) assert len(received) == 1 assert received[0].task_id == "t1" assert received[0].source_agent == "agent_a" async def test_handle_handoff_no_handler_does_nothing(self): manager = HandoffManager() handoff = HandoffMessage( source_agent="agent_a", target_agent="agent_b", task_id="t1", task_type="search", context={}, reason="test", ) # 不应报错 await manager._handle_handoff(handoff) async def test_handle_handoff_handler_error_is_caught(self): manager = HandoffManager() async def bad_handler(msg): raise ValueError("handler error") manager.register_handler("agent_b", bad_handler) handoff = HandoffMessage( source_agent="agent_a", target_agent="agent_b", task_id="t1", task_type="search", context={}, reason="test", ) # 不应抛出异常 await manager._handle_handoff(handoff) async def test_handle_handoff_multiple_handlers(self): manager = HandoffManager() results = [] async def handler1(msg): results.append("handler1") async def handler2(msg): results.append("handler2") manager.register_handler("agent_b", handler1) manager.register_handler("agent_b", handler2) handoff = HandoffMessage( source_agent="agent_a", target_agent="agent_b", task_id="t1", task_type="search", context={}, reason="test", ) await manager._handle_handoff(handoff) assert len(results) == 2 assert "handler1" in results assert "handler2" in results # ── HandoffManager Redis Pub/Sub 测试 ─────────────────────────── def _redis_available(): """检查 Redis 是否可用""" import os import redis url = os.environ.get("REDIS_URL", "redis://localhost:6381/0") try: r = redis.from_url(url) r.ping() r.close() return True except Exception: return False redis_available = _redis_available() @pytest.mark.redis class TestHandoffManagerRedisMode: """HandoffManager Redis Pub/Sub 测试（需要 Redis）""" @pytest.mark.skipif(not redis_available, reason="Redis not available") async def test_send_handoff_publishes_to_channel(self, redis_client, clean_redis): manager = HandoffManager(redis=redis_client) handoff = HandoffMessage( source_agent="agent_a", target_agent="agent_b", task_id="t1", task_type="search", context={"q": "hello"}, reason="delegation", ) await manager.send_handoff(handoff) # 验证消息发布到了正确的频道 pubsub = redis_client.pubsub() await pubsub.subscribe("agent:agent_b:handoff") # 等待订阅确认消息 msg = await asyncio.wait_for(pubsub.get_message(timeout=2.0), timeout=3.0) # 第一条消息是订阅确认，跳过 # 由于 publish 是 fire-and-forget，消息可能已经发送了 # 我们通过另一种方式验证：重新发送并监听 await manager.send_handoff(handoff) # 读取发布的消息 while True: msg = await asyncio.wait_for(pubsub.get_message(timeout=2.0), timeout=3.0) if msg and msg.get("type") == "message": data = json.loads(msg["data"]) assert data["source_agent"] == "agent_a" assert data["target_agent"] == "agent_b" assert data["task_id"] == "t1" assert data["reason"] == "delegation" break await pubsub.unsubscribe("agent:agent_b:handoff") @pytest.mark.skipif(not redis_available, reason="Redis not available") async def test_send_handoff_channel_format(self, redis_client, clean_redis): """验证 handoff 消息发送到 agent:{target_agent}:handoff 频道""" manager = HandoffManager(redis=redis_client) handoff = HandoffMessage( source_agent="planner", target_agent="executor", task_id="t2", task_type="execute", context={"plan": "step1"}, reason="execute plan", ) await manager.send_handoff(handoff) # 验证频道名格式 pubsub = redis_client.pubsub() await pubsub.subscribe("agent:executor:handoff") # 等待订阅确认 await asyncio.wait_for(pubsub.get_message(timeout=2.0), timeout=3.0) await manager.send_handoff(handoff) while True: msg = await asyncio.wait_for(pubsub.get_message(timeout=2.0), timeout=3.0) if msg and msg.get("type") == "message": data = json.loads(msg["data"]) assert data["target_agent"] == "executor" break await pubsub.unsubscribe("agent:executor:handoff") @pytest.mark.skipif(not redis_available, reason="Redis not available") async def test_different_agents_different_channels(self, redis_client, clean_redis): """不同 Agent 监听不同频道""" manager = HandoffManager(redis=redis_client) handoff_b = HandoffMessage( source_agent="a", target_agent="b", task_id="t3", task_type="search", context={}, reason="to b", ) handoff_c = HandoffMessage( source_agent="a", target_agent="c", task_id="t4", task_type="search", context={}, reason="to c", ) # 订阅 agent_b 的频道 pubsub_b = redis_client.pubsub() await pubsub_b.subscribe("agent:b:handoff") # 订阅 agent_c 的频道 pubsub_c = redis_client.pubsub() await pubsub_c.subscribe("agent:c:handoff") # 等待订阅确认 await asyncio.wait_for(pubsub_b.get_message(timeout=2.0), timeout=3.0) await asyncio.wait_for(pubsub_c.get_message(timeout=2.0), timeout=3.0) # 发送 handoff await manager.send_handoff(handoff_b) await manager.send_handoff(handoff_c) # 验证 b 收到自己的消息 while True: msg = await asyncio.wait_for(pubsub_b.get_message(timeout=2.0), timeout=3.0) if msg and msg.get("type") == "message": data = json.loads(msg["data"]) assert data["target_agent"] == "b" break # 验证 c 收到自己的消息 while True: msg = await asyncio.wait_for(pubsub_c.get_message(timeout=2.0), timeout=3.0) if msg and msg.get("type") == "message": data = json.loads(msg["data"]) assert data["target_agent"] == "c" break await pubsub_b.unsubscribe("agent:b:handoff") await pubsub_c.unsubscribe("agent:c:handoff") @pytest.mark.skipif(not redis_available, reason="Redis not available") async def test_listen_for_handoffs_receives_and_handles(self, redis_client, clean_redis): """listen_for_handoffs 接收消息并调用 handler""" manager = HandoffManager(redis=redis_client) received = [] async def handler(msg): received.append(msg) manager.register_handler("agent_b", handler) # 启动监听任务 listen_task = asyncio.create_task( manager.listen_for_handoffs("agent_b") ) # 等待订阅建立 await asyncio.sleep(0.5) # 发送 handoff handoff = HandoffMessage( source_agent="agent_a", target_agent="agent_b", task_id="t5", task_type="search", context={"q": "test"}, reason="delegation", ) await manager.send_handoff(handoff) # 等待处理 await asyncio.sleep(1.0) # 取消监听任务 listen_task.cancel() try: await listen_task except asyncio.CancelledError: pass assert len(received) == 1 assert received[0].task_id == "t5" assert received[0].source_agent == "agent_a" assert received[0].target_agent == "agent_b" assert received[0].context == {"q": "test"} assert received[0].reason == "delegation" @pytest.mark.skipif(not redis_available, reason="Redis not available") async def test_handoff_message_contains_all_fields(self, redis_client, clean_redis): """验证 handoff 消息包含 source_agent, target_agent, context, reason""" manager = HandoffManager(redis=redis_client) handoff = HandoffMessage( source_agent="researcher", target_agent="writer", task_id="t6", task_type="compose", context={"research": "findings", "style": "formal"}, reason="needs writing expertise", ) await manager.send_handoff(handoff) pubsub = redis_client.pubsub() await pubsub.subscribe("agent:writer:handoff") # 等待订阅确认 await asyncio.wait_for(pubsub.get_message(timeout=2.0), timeout=3.0) await manager.send_handoff(handoff) while True: msg = await asyncio.wait_for(pubsub.get_message(timeout=2.0), timeout=3.0) if msg and msg.get("type") == "message": data = json.loads(msg["data"]) assert data["source_agent"] == "researcher" assert data["target_agent"] == "writer" assert data["context"] == {"research": "findings", "style": "formal"} assert data["reason"] == "needs writing expertise" assert data["task_id"] == "t6" assert data["task_type"] == "compose" assert "created_at" in data break await pubsub.unsubscribe("agent:writer:handoff")