learningcircuit--local-deep-research/tests/retriever_integration/test_search_engine_retriever.py

"""Tests for the RetrieverSearchEngine."""

import pytest
from langchain_core.retrievers import BaseRetriever, Document

from local_deep_research.web_search_engines.engines.search_engine_retriever import (
    RetrieverSearchEngine,
)


class MockRetriever(BaseRetriever):
    """Mock retriever that returns predefined documents."""

    documents: list = []
    call_count: int = 0
    last_query: str = ""

    class Config:
        """Pydantic config."""

        arbitrary_types_allowed = True

    def _get_relevant_documents(self, query: str):
        self.call_count += 1
        self.last_query = query
        return self.documents

    async def _aget_relevant_documents(self, query: str):
        self.call_count += 1
        self.last_query = query
        return self.documents


class TestRetrieverSearchEngine:
    """Test the RetrieverSearchEngine functionality."""

    def test_basic_search(self):
        """Test basic search functionality."""
        # Create mock documents
        docs = [
            Document(
                page_content="First document content",
                metadata={"title": "Doc 1", "source": "test_source"},
            ),
            Document(
                page_content="Second document content",
                metadata={"title": "Doc 2", "source": "test_source"},
            ),
        ]

        retriever = MockRetriever(documents=docs)
        engine = RetrieverSearchEngine(
            retriever=retriever, name="test", max_results=10
        )

        results = engine.run("test query")

        assert len(results) == 2
        assert retriever.call_count == 1
        assert retriever.last_query == "test query"

        # Check result format
        assert results[0]["title"] == "Doc 1"
        assert results[0]["snippet"] == "First document content"
        assert results[0]["full_content"] == "First document content"
        assert results[0]["source"] == "test"

    def test_max_results_limit(self):
        """Test that max_results is respected."""
        # Create more documents than max_results
        docs = [Document(page_content=f"Document {i}") for i in range(20)]

        retriever = MockRetriever(documents=docs)
        engine = RetrieverSearchEngine(retriever=retriever, max_results=5)

        results = engine.run("test")
        assert len(results) == 5  # Should be limited to max_results

    def test_empty_results(self):
        """Test handling of empty results."""
        retriever = MockRetriever(documents=[])  # No documents
        engine = RetrieverSearchEngine(retriever=retriever)

        results = engine.run("test")
        assert results == []

    def test_document_conversion(self):
        """Test conversion of LangChain documents to LDR format."""
        doc = Document(
            page_content="Test content with more than 500 characters. " * 50,
            metadata={
                "title": "Test Title",
                "source": "http://example.com",
                "author": "Test Author",
                "date": "2024-01-01",
                "score": 0.95,
                "custom_field": "custom_value",
            },
        )

        retriever = MockRetriever(documents=[doc])
        engine = RetrieverSearchEngine(retriever=retriever, name="test_engine")

        results = engine.run("test")
        result = results[0]

        assert result["title"] == "Test Title"
        assert result["url"] == "http://example.com"
        assert len(result["snippet"]) == 500  # Should be truncated
        assert result["full_content"] == doc.page_content
        assert result["author"] == "Test Author"
        assert result["date"] == "2024-01-01"
        assert result["score"] == 0.95
        assert result["metadata"]["custom_field"] == "custom_value"
        assert result["source"] == "test_engine"
        assert result["retriever_type"] == "MockRetriever"

    def test_missing_metadata(self):
        """Test handling of documents with missing metadata."""
        doc = Document(
            page_content="Content without metadata",
            metadata={},  # Empty metadata
        )

        retriever = MockRetriever(documents=[doc])
        engine = RetrieverSearchEngine(retriever=retriever, name="test")

        results = engine.run("test")
        result = results[0]

        assert result["title"] == "Document 1"  # Default title
        assert result["url"] == "retriever://test/doc_0"  # Default URL
        assert result["author"] == ""
        assert result["date"] == ""
        assert result["score"] == 1.0  # Default score

    @pytest.mark.asyncio
    async def test_async_search(self):
        """Test async search functionality."""
        docs = [
            Document(
                page_content="Async document", metadata={"title": "Async Doc"}
            )
        ]

        retriever = MockRetriever(documents=docs)
        engine = RetrieverSearchEngine(retriever=retriever)

        results = await engine.arun("async test")

        assert len(results) == 1
        assert results[0]["title"] == "Async Doc"
        assert retriever.call_count == 1

    @pytest.mark.asyncio
    async def test_async_fallback(self):
        """Test fallback to sync when async not supported."""

        # Create a retriever without async support
        class SyncOnlyRetriever(BaseRetriever):
            class Config:
                arbitrary_types_allowed = True

            def _get_relevant_documents(self, query: str):
                return [Document(page_content="Sync only")]

        retriever = SyncOnlyRetriever()
        engine = RetrieverSearchEngine(retriever=retriever)

        # Should fall back to sync version
        results = await engine.arun("test")
        assert len(results) == 1
        assert results[0]["snippet"] == "Sync only"

    def test_error_handling(self):
        """Test error handling in search."""

        # Create a retriever that raises an error
        class ErrorRetriever(BaseRetriever):
            class Config:
                arbitrary_types_allowed = True

            def _get_relevant_documents(self, query: str):
                raise Exception("Test error")

            async def _aget_relevant_documents(self, query: str):
                raise Exception("Test error")

        retriever = ErrorRetriever()
        engine = RetrieverSearchEngine(retriever=retriever)

        # Should return empty list on error
        results = engine.run("test")
        assert results == []

    def test_none_page_content(self):
        """Test handling of documents with None page_content."""
        doc = Document(page_content="", metadata={"title": "Empty Doc"})

        retriever = MockRetriever(documents=[doc])
        engine = RetrieverSearchEngine(retriever=retriever)

        results = engine.run("test")
        assert results[0]["snippet"] == ""
        assert results[0]["full_content"] == ""