"""OrganizationContext 与 AgentDiscovery 单元测试""" import pytest from agentkit.org.context import AgentProfile, OrganizationContext from agentkit.org.discovery import AgentDiscovery from agentkit.skills.base import Skill, SkillConfig from agentkit.skills.registry import SkillRegistry # ---- Fixtures ---- @pytest.fixture def org_context(): return OrganizationContext() @pytest.fixture def profile_rag(): return AgentProfile( name="rag_agent", agent_type="react", capabilities=["rag", "search"], skills=["rag_skill"], execution_mode="react", model="gpt-4", ) @pytest.fixture def profile_terminal(): return AgentProfile( name="terminal_agent", agent_type="react", capabilities=["terminal", "shell"], skills=["terminal_skill"], execution_mode="react", model="gpt-4", ) @pytest.fixture def profile_coder(): return AgentProfile( name="coder_agent", agent_type="rewoo", capabilities=["rag", "terminal", "code_gen"], skills=["coder_skill"], execution_mode="rewoo", model="claude-3", max_concurrency=3, ) # ---- OrganizationContext: 注册与注销 ---- class TestOrganizationContextRegister: """注册与注销 Agent 档案""" def test_register_agent(self, org_context, profile_rag): org_context.register_agent(profile_rag) assert org_context.get_agent_profile("rag_agent") is profile_rag def test_unregister_agent(self, org_context, profile_rag): org_context.register_agent(profile_rag) org_context.unregister_agent("rag_agent") assert org_context.get_agent_profile("rag_agent") is None def test_unregister_nonexistent_no_error(self, org_context): org_context.unregister_agent("nonexistent") # should not raise def test_register_overwrites_existing(self, org_context, profile_rag): org_context.register_agent(profile_rag) updated = AgentProfile( name="rag_agent", agent_type="react", capabilities=["rag", "search", "summarize"], skills=["rag_skill"], ) org_context.register_agent(updated) profile = org_context.get_agent_profile("rag_agent") assert profile is updated assert "summarize" in profile.capabilities def test_list_agents(self, org_context, profile_rag, profile_terminal): org_context.register_agent(profile_rag) org_context.register_agent(profile_terminal) agents = org_context.list_agents() assert len(agents) == 2 names = {a.name for a in agents} assert names == {"rag_agent", "terminal_agent"} def test_list_agents_empty(self, org_context): assert org_context.list_agents() == [] # ---- OrganizationContext: 能力查找 ---- class TestOrganizationContextFind: """find_best_agent() 测试""" def test_find_by_required_capabilities(self, org_context, profile_rag, profile_terminal): org_context.register_agent(profile_rag) org_context.register_agent(profile_terminal) result = org_context.find_best_agent(["rag"]) assert result is not None assert result.name == "rag_agent" def test_find_exact_capability_match(self, org_context, profile_rag, profile_coder): org_context.register_agent(profile_rag) org_context.register_agent(profile_coder) # 两者都有 rag，但 coder 还有 terminal result = org_context.find_best_agent(["rag", "terminal"]) assert result is not None assert result.name == "coder_agent" def test_find_no_match_returns_none(self, org_context, profile_rag): org_context.register_agent(profile_rag) result = org_context.find_best_agent(["nonexistent_capability"]) assert result is None def test_find_excluded_agents_skipped(self, org_context, profile_rag, profile_coder): org_context.register_agent(profile_rag) org_context.register_agent(profile_coder) result = org_context.find_best_agent(["rag"], exclude=["coder_agent"]) assert result is not None assert result.name == "rag_agent" def test_find_unavailable_agents_skipped(self, org_context, profile_rag, profile_coder): org_context.register_agent(profile_rag) org_context.register_agent(profile_coder) org_context.set_availability("coder_agent", False) result = org_context.find_best_agent(["rag", "terminal"]) assert result is None # coder is unavailable, rag doesn't have terminal def test_find_best_agent_with_load_balancing(self, org_context): low_load = AgentProfile( name="low_load_agent", agent_type="react", capabilities=["rag"], skills=["rag_skill"], current_load=0, ) high_load = AgentProfile( name="high_load_agent", agent_type="react", capabilities=["rag"], skills=["rag_skill"], current_load=5, ) org_context.register_agent(low_load) org_context.register_agent(high_load) result = org_context.find_best_agent(["rag"]) assert result is not None assert result.name == "low_load_agent" def test_find_capability_case_insensitive(self, org_context, profile_rag): org_context.register_agent(profile_rag) result = org_context.find_best_agent(["RAG"]) assert result is not None assert result.name == "rag_agent" # ---- OrganizationContext: 负载与可用性 ---- class TestOrganizationContextLoadAvailability: """update_load() 和 set_availability() 测试""" def test_update_load_increase(self, org_context, profile_rag): org_context.register_agent(profile_rag) org_context.update_load("rag_agent", 3) assert org_context.get_agent_profile("rag_agent").current_load == 3 def test_update_load_decrease(self, org_context, profile_rag): org_context.register_agent(profile_rag) org_context.update_load("rag_agent", 5) org_context.update_load("rag_agent", -2) assert org_context.get_agent_profile("rag_agent").current_load == 3 def test_update_load_never_below_zero(self, org_context, profile_rag): org_context.register_agent(profile_rag) org_context.update_load("rag_agent", -10) assert org_context.get_agent_profile("rag_agent").current_load == 0 def test_update_load_nonexistent_no_error(self, org_context): org_context.update_load("nonexistent", 1) # should not raise def test_set_availability(self, org_context, profile_rag): org_context.register_agent(profile_rag) org_context.set_availability("rag_agent", False) assert org_context.get_agent_profile("rag_agent").availability is False org_context.set_availability("rag_agent", True) assert org_context.get_agent_profile("rag_agent").availability is True def test_set_availability_nonexistent_no_error(self, org_context): org_context.set_availability("nonexistent", False) # should not raise # ---- OrganizationContext: from_agent_pool ---- class TestOrganizationContextFromPool: """from_agent_pool() 测试""" def test_from_agent_pool_builds_context(self): """从 AgentPool + SkillRegistry 构建 OrganizationContext""" skill_registry = SkillRegistry() skill_config = SkillConfig( name="my_skill", agent_type="react", capabilities=["rag", "search"], execution_mode="react", llm={"model": "gpt-4"}, max_concurrency=2, prompt={"identity": "Test"}, ) skill = Skill(config=skill_config) skill_registry.register(skill) # Mock agent_pool class FakeAgentPool: def list_agents(self): return [{"name": "my_skill", "agent_type": "react"}] ctx = OrganizationContext.from_agent_pool(FakeAgentPool(), skill_registry) profile = ctx.get_agent_profile("my_skill") assert profile is not None assert profile.agent_type == "react" assert "rag" in profile.capabilities assert "search" in profile.capabilities assert profile.execution_mode == "react" assert profile.model == "gpt-4" assert profile.max_concurrency == 2 def test_from_agent_pool_none_graceful(self): """agent_pool 或 skill_registry 为 None 时返回空上下文""" ctx = OrganizationContext.from_agent_pool(None, SkillRegistry()) assert ctx.list_agents() == [] class FakePool: def list_agents(self): return [] ctx = OrganizationContext.from_agent_pool(FakePool(), None) assert ctx.list_agents() == [] def test_from_agent_pool_agent_not_in_registry(self): """Agent 不在 skill_registry 中时使用默认值""" skill_registry = SkillRegistry() class FakeAgentPool: def list_agents(self): return [{"name": "unknown_agent", "agent_type": "direct"}] ctx = OrganizationContext.from_agent_pool(FakeAgentPool(), skill_registry) profile = ctx.get_agent_profile("unknown_agent") assert profile is not None assert profile.agent_type == "direct" assert profile.capabilities == [] assert profile.execution_mode == "react" # default assert profile.model == "default" # ---- AgentDiscovery ---- class TestAgentDiscoveryByCapability: """discover_by_capability() 测试""" def test_discover_by_capability(self, org_context, profile_rag, profile_coder): org_context.register_agent(profile_rag) org_context.register_agent(profile_coder) discovery = AgentDiscovery(org_context) result = discovery.discover_by_capability(["rag"]) names = {p.name for p in result} assert names == {"rag_agent", "coder_agent"} def test_discover_by_capability_no_match(self, org_context, profile_rag): org_context.register_agent(profile_rag) discovery = AgentDiscovery(org_context) result = discovery.discover_by_capability(["nonexistent"]) assert result == [] class TestAgentDiscoveryByMode: """discover_by_execution_mode() 测试""" def test_discover_by_execution_mode(self, org_context, profile_rag, profile_coder): org_context.register_agent(profile_rag) org_context.register_agent(profile_coder) discovery = AgentDiscovery(org_context) result = discovery.discover_by_execution_mode("rewoo") assert len(result) == 1 assert result[0].name == "coder_agent" def test_discover_by_execution_mode_no_match(self, org_context, profile_rag): org_context.register_agent(profile_rag) discovery = AgentDiscovery(org_context) result = discovery.discover_by_execution_mode("plan_exec") assert result == [] class TestAgentDiscoveryAvailable: """discover_available() 测试""" def test_discover_available(self, org_context, profile_rag, profile_coder): org_context.register_agent(profile_rag) org_context.register_agent(profile_coder) org_context.set_availability("coder_agent", False) discovery = AgentDiscovery(org_context) result = discovery.discover_available() names = {p.name for p in result} assert names == {"rag_agent"} class TestAgentDiscoveryRecommend: """recommend_agent() 测试""" def test_recommend_with_preferred_mode(self, org_context, profile_rag, profile_coder): org_context.register_agent(profile_rag) org_context.register_agent(profile_coder) discovery = AgentDiscovery(org_context) result = discovery.recommend_agent(["rag"], preferred_mode="rewoo") assert result is not None assert result.name == "coder_agent" def test_recommend_without_preferred_mode(self, org_context, profile_rag, profile_coder): org_context.register_agent(profile_rag) org_context.register_agent(profile_coder) discovery = AgentDiscovery(org_context) result = discovery.recommend_agent(["rag"]) assert result is not None # Both have rag, should pick lower load assert result.current_load == 0 def test_recommend_fallback_when_no_capability_match(self, org_context, profile_rag): org_context.register_agent(profile_rag) discovery = AgentDiscovery(org_context) result = discovery.recommend_agent(["nonexistent"]) # Falls back to any available agent assert result is not None assert result.name == "rag_agent" def test_recommend_returns_none_when_no_available(self, org_context, profile_rag): org_context.register_agent(profile_rag) org_context.set_availability("rag_agent", False) discovery = AgentDiscovery(org_context) result = discovery.recommend_agent(["rag"]) assert result is None def test_recommend_preferred_mode_no_match_uses_any_match(self, org_context, profile_rag): org_context.register_agent(profile_rag) discovery = AgentDiscovery(org_context) # rag_agent has react mode, but we prefer plan_exec result = discovery.recommend_agent(["rag"], preferred_mode="plan_exec") # No plan_exec match, but still has capability match assert result is not None assert result.name == "rag_agent"