tatlock/tests/agents/test_coordination.py

"""
Tests for multi-agent coordination engine.
"""

import pytest
from unittest.mock import AsyncMock, MagicMock, patch

from src.agents.coordination import (
    CoordinationEngine,
    get_coordination_engine,
    delegate_to_librarian,
)
from src.agents.protocol import (
    AgentResponse,
    AgentUnavailableError,
    DelegationIntent,
    DelegationReason,
)


@pytest.fixture
def coordination_engine():
    """Create a fresh coordination engine for testing."""
    return CoordinationEngine()


@pytest.fixture
def mock_registry():
    """Mock the household registry."""
    with patch("src.agents.coordination.get_household_registry") as mock:
        registry = MagicMock()
        mock.return_value = registry
        yield registry


@pytest.fixture
def librarian_intent():
    """Create a standard librarian delegation intent."""
    return DelegationIntent(
        target_agent="librarian",
        task="Find information about Docker networking",
        reason=DelegationReason.DOMAIN_EXPERTISE,
        expected_outcome="Documentation and examples",
    )


@pytest.mark.unit
class TestCoordinationEngine:
    """Tests for CoordinationEngine class."""

    def test_initialization(self, coordination_engine):
        """Test engine initializes correctly."""
        assert coordination_engine is not None
        assert coordination_engine.registry is not None

    def test_get_available_agents_empty(self, mock_registry):
        """Test getting available agents when none have agents."""
        mock_registry.list_members.return_value = ["tatlock_core"]
        mock_member = MagicMock()
        mock_member.agent = None  # No agent
        mock_registry.get_member.return_value = mock_member

        engine = CoordinationEngine()
        available = engine.get_available_agents()

        assert available == []

    def test_get_available_agents_with_librarian(self, mock_registry):
        """Test getting available agents with librarian registered."""
        mock_registry.list_members.return_value = ["tatlock_core", "librarian"]

        # tatlock_core has no agent
        core_member = MagicMock()
        core_member.agent = None

        # librarian has an agent
        librarian_member = MagicMock()
        librarian_member.agent = MagicMock()

        def get_member_side_effect(name):
            if name == "tatlock_core":
                return core_member
            elif name == "librarian":
                return librarian_member
            return None

        mock_registry.get_member.side_effect = get_member_side_effect

        engine = CoordinationEngine()
        available = engine.get_available_agents()

        assert "librarian" in available
        assert "tatlock_core" not in available

    def test_can_delegate_to_unknown_agent(self, mock_registry):
        """Test checking delegation to unknown agent."""
        mock_registry.get_member.return_value = None

        engine = CoordinationEngine()

        assert engine.can_delegate_to("unknown_agent") is False

    def test_can_delegate_to_librarian(self, mock_registry):
        """Test checking delegation to librarian."""
        mock_member = MagicMock()
        mock_member.agent = MagicMock()  # Has an agent
        mock_registry.get_member.return_value = mock_member

        engine = CoordinationEngine()

        assert engine.can_delegate_to("librarian") is True


@pytest.mark.unit
class TestDelegationExecution:
    """Tests for delegation execution."""

    @pytest.mark.asyncio
    async def test_execute_delegation_unavailable_agent(
        self, mock_registry, librarian_intent
    ):
        """Test delegation fails for unavailable agent."""
        mock_registry.get_member.return_value = None

        engine = CoordinationEngine()

        with pytest.raises(AgentUnavailableError) as exc_info:
            await engine.execute_delegation(librarian_intent)

        assert "librarian" in str(exc_info.value)

    @pytest.mark.asyncio
    async def test_execute_delegation_success(
        self, mock_registry, librarian_intent
    ):
        """Test successful delegation execution."""
        # Setup mock member with agent
        mock_member = MagicMock()
        mock_member.agent = MagicMock()
        mock_registry.get_member.return_value = mock_member

        # Mock the executor
        with patch(
            "src.agents.coordination.AGENT_EXECUTORS",
            {"librarian": AsyncMock(return_value="Research results here")},
        ):
            engine = CoordinationEngine()
            response = await engine.execute_delegation(librarian_intent)

        assert response.success is True
        assert response.result == "Research results here"
        # Duration might be 0 for very fast mock execution
        assert response.duration_ms >= 0

    @pytest.mark.asyncio
    async def test_execute_delegation_error(
        self, mock_registry, librarian_intent
    ):
        """Test delegation handles executor errors."""
        mock_member = MagicMock()
        mock_member.agent = MagicMock()
        mock_registry.get_member.return_value = mock_member

        # Mock executor that raises
        async def failing_executor(**kwargs):
            raise ValueError("API connection failed")

        with patch(
            "src.agents.coordination.AGENT_EXECUTORS",
            {"librarian": failing_executor},
        ):
            engine = CoordinationEngine()
            response = await engine.execute_delegation(librarian_intent)

        assert response.success is False
        assert "API connection failed" in response.error_message


@pytest.mark.unit
class TestCoordinate:
    """Tests for multi-agent coordination."""

    @pytest.mark.asyncio
    async def test_coordinate_single_intent(self, mock_registry, librarian_intent):
        """Test coordinating a single delegation."""
        mock_member = MagicMock()
        mock_member.agent = MagicMock()
        mock_registry.get_member.return_value = mock_member

        with patch(
            "src.agents.coordination.AGENT_EXECUTORS",
            {"librarian": AsyncMock(return_value="Found docs")},
        ):
            engine = CoordinationEngine()
            result = await engine.coordinate([librarian_intent])

        assert result.final_response == "Found docs"
        assert "librarian" in result.agents_consulted
        # Duration might be 0 for very fast mock execution
        assert result.total_duration_ms >= 0

    @pytest.mark.asyncio
    async def test_coordinate_empty_intents(self, mock_registry):
        """Test coordinating with no intents."""
        engine = CoordinationEngine()
        result = await engine.coordinate([])

        assert result.final_response == ""
        assert result.agents_consulted == []

    @pytest.mark.asyncio
    async def test_coordinate_multiple_intents(self, mock_registry):
        """Test coordinating multiple delegations."""
        mock_member = MagicMock()
        mock_member.agent = MagicMock()
        mock_registry.get_member.return_value = mock_member

        intents = [
            DelegationIntent(
                target_agent="librarian",
                task="Task 1",
                reason=DelegationReason.DOMAIN_EXPERTISE,
                expected_outcome="Result 1",
                priority=1,
            ),
            DelegationIntent(
                target_agent="librarian",
                task="Task 2",
                reason=DelegationReason.DOMAIN_EXPERTISE,
                expected_outcome="Result 2",
                priority=2,
            ),
        ]

        call_count = 0

        async def mock_executor(**kwargs):
            nonlocal call_count
            call_count += 1
            return f"Result {call_count}"

        with patch(
            "src.agents.coordination.AGENT_EXECUTORS",
            {"librarian": mock_executor},
        ):
            engine = CoordinationEngine()
            result = await engine.coordinate(intents)

        # Both intents were executed (check agents_consulted count)
        assert len(result.agents_consulted) == 2
        # Current implementation replaces same-agent responses in dict
        # So final_response has the last result (or combined if different agents)
        assert len(result.final_response) > 0


@pytest.mark.unit
class TestDelegateToLibrarian:
    """Tests for convenience delegation function."""

    @pytest.mark.asyncio
    async def test_delegate_to_librarian(self, mock_registry):
        """Test the delegate_to_librarian helper."""
        mock_member = MagicMock()
        mock_member.agent = MagicMock()
        mock_registry.get_member.return_value = mock_member

        with patch(
            "src.agents.coordination.AGENT_EXECUTORS",
            {"librarian": AsyncMock(return_value="Wiki search results")},
        ):
            # Reset global engine
            with patch(
                "src.agents.coordination._coordination_engine",
                None,
            ):
                response = await delegate_to_librarian(
                    task="Search for Docker docs",
                    context="Setting up homelab",
                )

        assert response.success is True
        assert response.result == "Wiki search results"


@pytest.mark.unit
class TestGetCoordinationEngine:
    """Tests for engine singleton."""

    def test_get_coordination_engine_singleton(self):
        """Test engine is singleton."""
        with patch("src.agents.coordination._coordination_engine", None):
            engine1 = get_coordination_engine()
            engine2 = get_coordination_engine()

            # Should be same instance
            assert engine1 is engine2


@pytest.mark.unit
class TestDelegationStreaming:
    """Tests for streaming delegation."""

    @pytest.mark.asyncio
    async def test_execute_delegation_stream_unavailable(
        self, mock_registry, librarian_intent
    ):
        """Test streaming fails for unavailable agent."""
        engine = CoordinationEngine()

        # Change target to an agent that doesn't have a stream executor
        librarian_intent.target_agent = "nonexistent_agent"

        with pytest.raises(AgentUnavailableError):
            async for _ in engine.execute_delegation_stream(librarian_intent):
                pass

    @pytest.mark.asyncio
    async def test_execute_delegation_stream_success(
        self, mock_registry, librarian_intent
    ):
        """Test successful streaming delegation."""
        mock_member = MagicMock()
        mock_member.agent = MagicMock()
        mock_registry.get_member.return_value = mock_member

        async def mock_stream(**kwargs):
            yield "Hello "
            yield "world"

        with patch(
            "src.agents.coordination.AGENT_STREAM_EXECUTORS",
            {"librarian": mock_stream},
        ):
            engine = CoordinationEngine()
            chunks = []
            async for chunk in engine.execute_delegation_stream(librarian_intent):
                chunks.append(chunk)

        assert chunks == ["Hello ", "world"]