tests/unit/tools/test_search_handler.py

"""Unit tests for SearchHandler."""

from unittest.mock import AsyncMock

import pytest

from src.tools.search_handler import SearchHandler
from src.utils.exceptions import SearchError
from src.utils.models import Citation, Evidence


class TestSearchHandler:
    """Tests for SearchHandler."""

    @pytest.mark.asyncio
    async def test_execute_aggregates_results(self):
        """SearchHandler should aggregate results from all tools."""
        # Create mock tools
        mock_tool_1 = AsyncMock()
        mock_tool_1.name = "pubmed"
        mock_tool_1.search = AsyncMock(
            return_value=[
                Evidence(
                    content="Result 1",
                    citation=Citation(source="pubmed", title="T1", url="u1", date="2024"),
                )
            ]
        )

        mock_tool_2 = AsyncMock()
        mock_tool_2.name = "pubmed"  # Type system currently restricts to pubmed
        mock_tool_2.search = AsyncMock(return_value=[])

        handler = SearchHandler(tools=[mock_tool_1, mock_tool_2])
        result = await handler.execute("test query")

        assert result.total_found == 1
        assert "pubmed" in result.sources_searched
        assert len(result.errors) == 0

    @pytest.mark.asyncio
    async def test_execute_handles_tool_failure(self):
        """SearchHandler should continue if one tool fails."""
        mock_tool_ok = AsyncMock()
        mock_tool_ok.name = "pubmed"
        mock_tool_ok.search = AsyncMock(
            return_value=[
                Evidence(
                    content="Good result",
                    citation=Citation(source="pubmed", title="T", url="u", date="2024"),
                )
            ]
        )

        mock_tool_fail = AsyncMock()
        mock_tool_fail.name = "pubmed"  # Mocking a second pubmed instance failing
        mock_tool_fail.search = AsyncMock(side_effect=SearchError("API down"))

        handler = SearchHandler(tools=[mock_tool_ok, mock_tool_fail])
        result = await handler.execute("test")

        assert result.total_found == 1
        assert "pubmed" in result.sources_searched
        assert len(result.errors) == 1
        # The error message format is "{tool.name}: {error!s}"
        assert "pubmed: API down" in result.errors[0]

    @pytest.mark.asyncio
    async def test_search_handler_pubmed_only(self):
        """SearchHandler should work with only PubMed tool."""
        # This is the specific test requested in Phase 9 spec
        from src.tools.pubmed import PubMedTool

        mock_pubmed = AsyncMock(spec=PubMedTool)
        mock_pubmed.name = "pubmed"
        mock_pubmed.search.return_value = []

        handler = SearchHandler(tools=[mock_pubmed], timeout=30.0)
        result = await handler.execute("metformin diabetes", max_results_per_tool=3)

        assert result.sources_searched == ["pubmed"]
        assert "web" not in result.sources_searched
        assert len(result.errors) == 0