tatlock/tests/agents/librarian/test_client_retry.py

"""
Tests for bounded retries, timeout wiring, and client reuse in
LibraryDeskClient, plus ModelRetry escalation from the read tools.
"""

from unittest.mock import AsyncMock, MagicMock, patch

import httpx
import pytest
from pydantic_ai import ModelRetry

from src.agents.librarian.client import (
    LibraryDeskClient,
    library_client_session,
)
from src.agents.librarian.tools import (
    create_wiki_page,
    hybrid_search,
    search_wiki,
)
from src.core.config import config


@pytest.fixture(autouse=True)
def _no_backoff(monkeypatch):
    """Skip the retry backoff sleep in tests."""
    monkeypatch.setattr("src.agents.librarian.client._RETRY_BACKOFF_SECONDS", 0)


def _ok_response(payload: dict) -> MagicMock:
    response = MagicMock()
    response.status_code = 200
    response.json.return_value = payload
    response.raise_for_status = MagicMock()
    return response


@pytest.fixture
def client_with_mock():
    client = LibraryDeskClient(base_url="http://test:8089", api_key="test-key")
    client._client = AsyncMock(spec=httpx.AsyncClient)
    return client


@pytest.mark.unit
class TestBoundedRetries:
    """2-attempt retry for GETs and read-only POST /query/*, /rag/search."""

    @pytest.mark.asyncio
    async def test_get_retries_once_on_transport_error(self, client_with_mock):
        mock_httpx = client_with_mock._client
        mock_httpx.get.side_effect = [
            httpx.ConnectError("Connection refused"),
            _ok_response({"results": []}),
        ]

        results = await client_with_mock.search_wiki("docker", user="u")

        assert results == []
        assert mock_httpx.get.call_count == 2

    @pytest.mark.asyncio
    async def test_get_gives_up_after_two_attempts(self, client_with_mock):
        mock_httpx = client_with_mock._client
        mock_httpx.get.side_effect = httpx.ConnectError("Connection refused")

        with pytest.raises(httpx.ConnectError):
            await client_with_mock.search_wiki("docker", user="u")

        assert mock_httpx.get.call_count == 2

    @pytest.mark.asyncio
    async def test_query_hybrid_retries_on_503(self, client_with_mock):
        mock_httpx = client_with_mock._client
        bad = MagicMock()
        bad.status_code = 503
        mock_httpx.post.side_effect = [
            bad,
            _ok_response({"results": [], "keywords": {}, "context": ""}),
        ]

        response = await client_with_mock.hybrid_search("docker", user="u")

        assert response.results == []
        assert mock_httpx.post.call_count == 2

    @pytest.mark.asyncio
    async def test_wiki_write_is_never_retried(self, client_with_mock):
        """POST /wiki/pages must not retry - it could duplicate pages."""
        mock_httpx = client_with_mock._client
        mock_httpx.post.side_effect = httpx.ConnectError("Connection refused")

        with pytest.raises(httpx.ConnectError):
            await client_with_mock.create_wiki_page(
                title="T", path="/t", content="c", user="u"
            )

        assert mock_httpx.post.call_count == 1

    @pytest.mark.asyncio
    async def test_smart_create_is_never_retried(self, client_with_mock):
        mock_httpx = client_with_mock._client
        mock_httpx.post.side_effect = httpx.ConnectError("Connection refused")

        with pytest.raises(httpx.ConnectError):
            await client_with_mock.smart_create_wiki_page(
                topic="T", tags=["x"], user="u"
            )

        assert mock_httpx.post.call_count == 1


@pytest.mark.unit
class TestTimeoutWiring:
    """LIBRARY_DESK_TIMEOUT config replaces the hardcoded 60s/30s."""

    def test_default_timeout_from_config(self):
        client = LibraryDeskClient()

        assert client.timeout == config.LIBRARY_DESK_TIMEOUT

    def test_explicit_timeout_wins(self):
        client = LibraryDeskClient(timeout=5)

        assert client.timeout == 5


@pytest.mark.unit
class TestClientReuse:
    """One shared HTTP connection per librarian run."""

    @pytest.mark.asyncio
    async def test_clients_share_connection_inside_session(self):
        async with library_client_session():
            async with LibraryDeskClient() as c1:
                http1 = c1._client
            # shared connection survives client exit
            assert http1 is not None
            assert not http1.is_closed

            async with LibraryDeskClient() as c2:
                assert c2._client is http1

        # session close tears the shared connection down
        assert http1.is_closed

    @pytest.mark.asyncio
    async def test_nested_sessions_are_noops(self):
        async with library_client_session():
            async with LibraryDeskClient() as c1:
                http1 = c1._client
            async with library_client_session():
                async with LibraryDeskClient() as c2:
                    assert c2._client is http1
            # inner session exit must not close the shared connection
            assert not http1.is_closed

    @pytest.mark.asyncio
    async def test_client_owns_connection_outside_session(self):
        async with LibraryDeskClient() as client:
            http_client = client._client

        assert http_client.is_closed

    @pytest.mark.asyncio
    async def test_custom_target_does_not_reuse_shared(self):
        async with library_client_session():
            async with LibraryDeskClient() as shared_client:
                shared_http = shared_client._client
            async with LibraryDeskClient(base_url="http://other:9999") as custom:
                assert custom._client is not shared_http


@pytest.mark.unit
class TestModelRetryEscalation:
    """Read tools raise ModelRetry on transient errors so Agent(retries=2) engages."""

    def _patched_client(self, mock_client):
        factory = MagicMock()
        factory.return_value.__aenter__ = AsyncMock(return_value=mock_client)
        factory.return_value.__aexit__ = AsyncMock(return_value=None)
        return patch("src.agents.librarian.tools.LibraryDeskClient", factory)

    @pytest.mark.asyncio
    async def test_read_tool_raises_model_retry_on_transport_error(self):
        mock_client = AsyncMock()
        mock_client.hybrid_search.side_effect = httpx.ConnectError(
            "Connection refused"
        )

        with self._patched_client(mock_client):
            with pytest.raises(ModelRetry):
                await hybrid_search("docker")

    @pytest.mark.asyncio
    async def test_read_tool_raises_model_retry_on_5xx(self):
        request = httpx.Request("GET", "http://test:8089/wiki/search")
        response = httpx.Response(502, request=request)
        mock_client = AsyncMock()
        mock_client.search_wiki.side_effect = httpx.HTTPStatusError(
            "bad gateway", request=request, response=response
        )

        with self._patched_client(mock_client):
            with pytest.raises(ModelRetry):
                await search_wiki("docker")

    @pytest.mark.asyncio
    async def test_read_tool_returns_safe_message_on_non_transient(self):
        mock_client = AsyncMock()
        mock_client.hybrid_search.side_effect = ValueError("bad parse")

        with self._patched_client(mock_client):
            result = await hybrid_search("docker")

        assert "unable" in result
        assert "bad parse" not in result

    @pytest.mark.asyncio
    async def test_write_tool_never_raises_model_retry(self):
        mock_client = AsyncMock()
        mock_client.create_wiki_page.side_effect = httpx.ConnectError(
            "Connection refused"
        )

        with self._patched_client(mock_client):
            result = await create_wiki_page(
                title="T", path="/t", content="c", tags=["x"]
            )

        assert "unable" in result
        assert "Connection refused" not in result