Timmy-time-dashboard/tests/timmy/test_agents_base.py

"""Tests for timmy.agents.base — BaseAgent and SubAgent.

Covers:
- Initialization and default values
- Tool registry integration
- Event bus connection and subscription
- run() with retry logic (transient + fatal errors)
- Event emission on successful run
- get_capabilities / get_status
- SubAgent.execute_task delegation
"""

from unittest.mock import AsyncMock, MagicMock, patch

import httpx
import pytest

# ── helpers ──────────────────────────────────────────────────────────────────


def _mock_settings(**overrides):
    """Create a settings mock with sensible defaults."""
    s = MagicMock()
    s.ollama_model = "qwen3:30b"
    s.ollama_url = "http://localhost:11434"
    s.ollama_num_ctx = 0
    s.telemetry_enabled = False
    for k, v in overrides.items():
        setattr(s, k, v)
    return s


def _make_agent_class():
    """Import after patches are in place."""
    from timmy.agents.base import SubAgent

    return SubAgent


def _make_base_class():
    from timmy.agents.base import BaseAgent

    return BaseAgent


# ── patch context ────────────────────────────────────────────────────────────

# All tests patch Agno's Agent so we never touch Ollama.
_AGENT_PATCH = "timmy.agents.base.Agent"
_OLLAMA_PATCH = "timmy.agents.base.Ollama"
_SETTINGS_PATCH = "timmy.agents.base.settings"
_REGISTRY_PATCH = "timmy.agents.base.tool_registry"


# ── Initialization ───────────────────────────────────────────────────────────


class TestBaseAgentInit:
    @patch(_REGISTRY_PATCH, None)
    @patch(_SETTINGS_PATCH, _mock_settings())
    @patch(_OLLAMA_PATCH)
    @patch(_AGENT_PATCH)
    def test_defaults(self, mock_agent_cls, mock_ollama):
        SubAgent = _make_agent_class()
        agent = SubAgent(
            agent_id="test-1",
            name="TestBot",
            role="tester",
            system_prompt="You are a test agent.",
        )
        assert agent.agent_id == "test-1"
        assert agent.name == "TestBot"
        assert agent.role == "tester"
        assert agent.tools == []
        assert agent.model == "qwen3:30b"
        assert agent.max_history == 10
        assert agent.event_bus is None

    @patch(_REGISTRY_PATCH, None)
    @patch(_SETTINGS_PATCH, _mock_settings())
    @patch(_OLLAMA_PATCH)
    @patch(_AGENT_PATCH)
    def test_custom_model(self, mock_agent_cls, mock_ollama):
        SubAgent = _make_agent_class()
        agent = SubAgent(
            agent_id="a",
            name="A",
            role="r",
            system_prompt="p",
            model="llama3:8b",
        )
        assert agent.model == "llama3:8b"

    @patch(_REGISTRY_PATCH, None)
    @patch(_SETTINGS_PATCH, _mock_settings())
    @patch(_OLLAMA_PATCH)
    @patch(_AGENT_PATCH)
    def test_custom_max_history(self, mock_agent_cls, mock_ollama):
        SubAgent = _make_agent_class()
        agent = SubAgent(agent_id="a", name="A", role="r", system_prompt="p", max_history=5)
        assert agent.max_history == 5

    @patch(_REGISTRY_PATCH, None)
    @patch(_SETTINGS_PATCH, _mock_settings())
    @patch(_OLLAMA_PATCH)
    @patch(_AGENT_PATCH)
    def test_tools_list_stored(self, mock_agent_cls, mock_ollama):
        SubAgent = _make_agent_class()
        agent = SubAgent(
            agent_id="a",
            name="A",
            role="r",
            system_prompt="p",
            tools=["calculator", "search"],
        )
        assert agent.tools == ["calculator", "search"]


# ── _create_agent internals ──────────────────────────────────────────────────


class TestCreateAgent:
    @patch(_REGISTRY_PATCH, None)
    @patch(_SETTINGS_PATCH, _mock_settings(ollama_num_ctx=4096))
    @patch(_OLLAMA_PATCH)
    @patch(_AGENT_PATCH)
    def test_num_ctx_passed_when_set(self, mock_agent_cls, mock_ollama):
        SubAgent = _make_agent_class()
        SubAgent(agent_id="a", name="A", role="r", system_prompt="p")
        # Ollama should have been called with options
        _, kwargs = mock_ollama.call_args
        assert kwargs.get("options") == {"num_ctx": 4096}

    @patch(_REGISTRY_PATCH, None)
    @patch(_SETTINGS_PATCH, _mock_settings(ollama_num_ctx=0))
    @patch(_OLLAMA_PATCH)
    @patch(_AGENT_PATCH)
    def test_num_ctx_omitted_when_zero(self, mock_agent_cls, mock_ollama):
        SubAgent = _make_agent_class()
        SubAgent(agent_id="a", name="A", role="r", system_prompt="p")
        _, kwargs = mock_ollama.call_args
        assert "options" not in kwargs

    @patch(_SETTINGS_PATCH, _mock_settings())
    @patch(_OLLAMA_PATCH)
    @patch(_AGENT_PATCH)
    def test_tool_registry_lookup(self, mock_agent_cls, mock_ollama):
        mock_registry = MagicMock()
        handler1 = MagicMock()
        handler2 = None  # Simulate missing tool
        mock_registry.get_handler.side_effect = [handler1, handler2]

        with patch(_REGISTRY_PATCH, mock_registry):
            SubAgent = _make_agent_class()
            SubAgent(
                agent_id="a",
                name="A",
                role="r",
                system_prompt="p",
                tools=["calc", "missing"],
            )

        assert mock_registry.get_handler.call_count == 2
        # Agent should have been created with just the one handler
        _, kwargs = mock_agent_cls.call_args
        assert kwargs["tools"] == [handler1]

    @patch(_SETTINGS_PATCH, _mock_settings())
    @patch(_OLLAMA_PATCH)
    @patch(_AGENT_PATCH)
    def test_no_tools_passes_none(self, mock_agent_cls, mock_ollama):
        with patch(_REGISTRY_PATCH, None):
            SubAgent = _make_agent_class()
            SubAgent(agent_id="a", name="A", role="r", system_prompt="p")
        _, kwargs = mock_agent_cls.call_args
        assert kwargs["tools"] is None


# ── Event bus ────────────────────────────────────────────────────────────────


class TestEventBus:
    @patch(_REGISTRY_PATCH, None)
    @patch(_SETTINGS_PATCH, _mock_settings())
    @patch(_OLLAMA_PATCH)
    @patch(_AGENT_PATCH)
    def test_connect_event_bus(self, mock_agent_cls, mock_ollama):
        SubAgent = _make_agent_class()
        agent = SubAgent(agent_id="bot-1", name="B", role="r", system_prompt="p")
        bus = MagicMock()
        bus.subscribe.return_value = lambda fn: fn  # decorator pattern

        agent.connect_event_bus(bus)

        assert agent.event_bus is bus
        assert bus.subscribe.call_count == 2
        # Check subscription patterns
        patterns = [call.args[0] for call in bus.subscribe.call_args_list]
        assert "agent.bot-1.*" in patterns
        assert "agent.task.assigned" in patterns


# ── run() retry logic ────────────────────────────────────────────────────────


class TestRun:
    @patch(_REGISTRY_PATCH, None)
    @patch(_SETTINGS_PATCH, _mock_settings())
    @patch(_OLLAMA_PATCH)
    @patch(_AGENT_PATCH)
    @pytest.mark.asyncio
    async def test_run_success(self, mock_agent_cls, mock_ollama):
        SubAgent = _make_agent_class()
        agent = SubAgent(agent_id="a", name="A", role="r", system_prompt="p")
        mock_result = MagicMock()
        mock_result.content = "Hello world"
        agent.agent.run.return_value = mock_result

        response = await agent.run("Hi")

        assert response == "Hello world"
        agent.agent.run.assert_called_once_with("Hi", stream=False)

    @patch(_REGISTRY_PATCH, None)
    @patch(_SETTINGS_PATCH, _mock_settings())
    @patch(_OLLAMA_PATCH)
    @patch(_AGENT_PATCH)
    @pytest.mark.asyncio
    async def test_run_result_without_content(self, mock_agent_cls, mock_ollama):
        """When result has no .content, fall back to str()."""
        SubAgent = _make_agent_class()
        agent = SubAgent(agent_id="a", name="A", role="r", system_prompt="p")
        agent.agent.run.return_value = "plain string"

        response = await agent.run("Hi")

        assert response == "plain string"

    @patch(_REGISTRY_PATCH, None)
    @patch(_SETTINGS_PATCH, _mock_settings())
    @patch(_OLLAMA_PATCH)
    @patch(_AGENT_PATCH)
    @pytest.mark.asyncio
    async def test_run_retries_transient_error(self, mock_agent_cls, mock_ollama):
        """Transient errors (ConnectError etc.) should be retried."""
        SubAgent = _make_agent_class()
        agent = SubAgent(agent_id="a", name="A", role="r", system_prompt="p")

        mock_result = MagicMock()
        mock_result.content = "recovered"
        agent.agent.run.side_effect = [
            httpx.ConnectError("refused"),
            mock_result,
        ]

        with patch("asyncio.sleep", new_callable=AsyncMock):
            response = await agent.run("Hi")

        assert response == "recovered"
        assert agent.agent.run.call_count == 2

    @patch(_REGISTRY_PATCH, None)
    @patch(_SETTINGS_PATCH, _mock_settings())
    @patch(_OLLAMA_PATCH)
    @patch(_AGENT_PATCH)
    @pytest.mark.asyncio
    async def test_run_retries_read_timeout(self, mock_agent_cls, mock_ollama):
        """ReadTimeout (GPU contention) should be retried."""
        SubAgent = _make_agent_class()
        agent = SubAgent(agent_id="a", name="A", role="r", system_prompt="p")

        mock_result = MagicMock()
        mock_result.content = "ok"
        agent.agent.run.side_effect = [
            httpx.ReadTimeout("timeout"),
            mock_result,
        ]

        with patch("asyncio.sleep", new_callable=AsyncMock):
            response = await agent.run("Hi")

        assert response == "ok"

    @patch(_REGISTRY_PATCH, None)
    @patch(_SETTINGS_PATCH, _mock_settings())
    @patch(_OLLAMA_PATCH)
    @patch(_AGENT_PATCH)
    @pytest.mark.asyncio
    async def test_run_exhausts_retries_transient(self, mock_agent_cls, mock_ollama):
        """After 3 transient failures, should raise."""
        SubAgent = _make_agent_class()
        agent = SubAgent(agent_id="a", name="A", role="r", system_prompt="p")

        agent.agent.run.side_effect = httpx.ConnectError("down")

        with patch("asyncio.sleep", new_callable=AsyncMock):
            with pytest.raises(httpx.ConnectError):
                await agent.run("Hi")

        assert agent.agent.run.call_count == 3

    @patch(_REGISTRY_PATCH, None)
    @patch(_SETTINGS_PATCH, _mock_settings())
    @patch(_OLLAMA_PATCH)
    @patch(_AGENT_PATCH)
    @pytest.mark.asyncio
    async def test_run_retries_non_transient_error(self, mock_agent_cls, mock_ollama):
        """Non-transient errors also get retried (with different backoff)."""
        SubAgent = _make_agent_class()
        agent = SubAgent(agent_id="a", name="A", role="r", system_prompt="p")

        agent.agent.run.side_effect = ValueError("bad input")

        with patch("asyncio.sleep", new_callable=AsyncMock):
            with pytest.raises(ValueError, match="bad input"):
                await agent.run("Hi")

        assert agent.agent.run.call_count == 3

    @patch(_REGISTRY_PATCH, None)
    @patch(_SETTINGS_PATCH, _mock_settings())
    @patch(_OLLAMA_PATCH)
    @patch(_AGENT_PATCH)
    @pytest.mark.asyncio
    async def test_run_emits_event_on_success(self, mock_agent_cls, mock_ollama):
        """Successful run should publish response event to bus."""
        SubAgent = _make_agent_class()
        agent = SubAgent(agent_id="bot-1", name="B", role="r", system_prompt="p")

        mock_bus = AsyncMock()
        agent.event_bus = mock_bus

        mock_result = MagicMock()
        mock_result.content = "answer"
        agent.agent.run.return_value = mock_result

        await agent.run("question")

        mock_bus.publish.assert_called_once()
        event = mock_bus.publish.call_args[0][0]
        assert event.type == "agent.bot-1.response"
        assert event.data["input"] == "question"
        assert event.data["output"] == "answer"

    @patch(_REGISTRY_PATCH, None)
    @patch(_SETTINGS_PATCH, _mock_settings())
    @patch(_OLLAMA_PATCH)
    @patch(_AGENT_PATCH)
    @pytest.mark.asyncio
    async def test_run_no_event_without_bus(self, mock_agent_cls, mock_ollama):
        """No bus connected = no event emitted (no crash)."""
        SubAgent = _make_agent_class()
        agent = SubAgent(agent_id="a", name="A", role="r", system_prompt="p")

        mock_result = MagicMock()
        mock_result.content = "ok"
        agent.agent.run.return_value = mock_result

        # Should not raise
        response = await agent.run("Hi")
        assert response == "ok"


# ── get_capabilities / get_status ────────────────────────────────────────────


class TestStatusAndCapabilities:
    @patch(_REGISTRY_PATCH, None)
    @patch(_SETTINGS_PATCH, _mock_settings())
    @patch(_OLLAMA_PATCH)
    @patch(_AGENT_PATCH)
    def test_get_capabilities(self, mock_agent_cls, mock_ollama):
        SubAgent = _make_agent_class()
        agent = SubAgent(agent_id="a", name="A", role="r", system_prompt="p", tools=["t1", "t2"])
        assert agent.get_capabilities() == ["t1", "t2"]

    @patch(_REGISTRY_PATCH, None)
    @patch(_SETTINGS_PATCH, _mock_settings())
    @patch(_OLLAMA_PATCH)
    @patch(_AGENT_PATCH)
    def test_get_status(self, mock_agent_cls, mock_ollama):
        SubAgent = _make_agent_class()
        agent = SubAgent(
            agent_id="bot-1",
            name="TestBot",
            role="assistant",
            system_prompt="p",
            tools=["calc"],
        )
        status = agent.get_status()
        assert status == {
            "agent_id": "bot-1",
            "name": "TestBot",
            "role": "assistant",
            "model": "qwen3:30b",
            "status": "ready",
            "tools": ["calc"],
        }


# ── SubAgent.execute_task ────────────────────────────────────────────────────


class TestSubAgentExecuteTask:
    @patch(_REGISTRY_PATCH, None)
    @patch(_SETTINGS_PATCH, _mock_settings())
    @patch(_OLLAMA_PATCH)
    @patch(_AGENT_PATCH)
    @pytest.mark.asyncio
    async def test_execute_task_delegates_to_run(self, mock_agent_cls, mock_ollama):
        SubAgent = _make_agent_class()
        agent = SubAgent(agent_id="bot-1", name="B", role="r", system_prompt="p")

        mock_result = MagicMock()
        mock_result.content = "task done"
        agent.agent.run.return_value = mock_result

        result = await agent.execute_task("t-1", "do the thing", {"extra": True})

        assert result == {
            "task_id": "t-1",
            "agent": "bot-1",
            "result": "task done",
            "status": "completed",
        }


# ── Task assignment handler ──────────────────────────────────────────────────


class TestTaskAssignment:
    @patch(_REGISTRY_PATCH, None)
    @patch(_SETTINGS_PATCH, _mock_settings())
    @patch(_OLLAMA_PATCH)
    @patch(_AGENT_PATCH)
    @pytest.mark.asyncio
    async def test_handles_assigned_task(self, mock_agent_cls, mock_ollama):
        """Agent should process tasks assigned to it."""
        SubAgent = _make_agent_class()

        agent = SubAgent(agent_id="bot-1", name="B", role="r", system_prompt="p")

        mock_result = MagicMock()
        mock_result.content = "done"
        agent.agent.run.return_value = mock_result

        from infrastructure.events.bus import Event

        event = Event(
            type="agent.task.assigned",
            source="coordinator",
            data={
                "agent_id": "bot-1",
                "task_id": "task-42",
                "description": "Fix the bug",
            },
        )

        await agent._handle_task_assignment(event)
        agent.agent.run.assert_called_once_with("Fix the bug", stream=False)

    @patch(_REGISTRY_PATCH, None)
    @patch(_SETTINGS_PATCH, _mock_settings())
    @patch(_OLLAMA_PATCH)
    @patch(_AGENT_PATCH)
    @pytest.mark.asyncio
    async def test_ignores_task_for_other_agent(self, mock_agent_cls, mock_ollama):
        """Agent should ignore tasks assigned to someone else."""
        SubAgent = _make_agent_class()
        agent = SubAgent(agent_id="bot-1", name="B", role="r", system_prompt="p")

        from infrastructure.events.bus import Event

        event = Event(
            type="agent.task.assigned",
            source="coordinator",
            data={
                "agent_id": "bot-2",
                "task_id": "task-99",
                "description": "Not my job",
            },
        )

        await agent._handle_task_assignment(event)
        agent.agent.run.assert_not_called()