"""
Tests for the AgentRegistry.
"""

import pytest
from dcos.core.registry import AgentRegistry, AgentRecord


class TestAgentRegistry:
    def test_register_and_get(self):
        registry = AgentRegistry()
        record = AgentRecord(agent_id="a1", name="test-agent", role="planner")
        registry.register(record)
        retrieved = registry.get("a1")
        assert retrieved is not None
        assert retrieved.name == "test-agent"

    def test_unregister(self):
        registry = AgentRegistry()
        record = AgentRecord(agent_id="a2", name="removable", role="researcher")
        registry.register(record)
        removed = registry.unregister("a2")
        assert removed is not None
        assert registry.get("a2") is None

    def test_find_by_role(self):
        registry = AgentRegistry()
        registry.register(AgentRecord(agent_id="a1", name="planner-1", role="planner"))
        registry.register(AgentRecord(agent_id="a2", name="planner-2", role="planner"))
        registry.register(AgentRecord(agent_id="a3", name="researcher", role="researcher"))
        planners = registry.find_by_role("planner")
        assert len(planners) == 2

    def test_find_by_capability(self):
        registry = AgentRegistry()
        registry.register(AgentRecord(agent_id="a1", name="capable", role="planner", capabilities=["ml", "nlp"]))
        registry.register(AgentRecord(agent_id="a2", name="other", role="researcher", capabilities=["search"]))
        ml_agents = registry.find_by_capability("ml")
        assert len(ml_agents) == 1
        assert ml_agents[0].agent_id == "a1"

    def test_count(self):
        registry = AgentRegistry()
        assert registry.count() == 0
        registry.register(AgentRecord(agent_id="a1", name="only", role="planner"))
        assert registry.count() == 1