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

"""Tests for retriever integration with search engine factory."""

import pytest
from langchain_core.retrievers import BaseRetriever, Document

from local_deep_research.web_search_engines.engines.search_engine_retriever import (
    RetrieverSearchEngine,
)
from local_deep_research.web_search_engines.retriever_registry import (
    retriever_registry,
)
from local_deep_research.web_search_engines.search_engine_factory import (
    create_search_engine,
)

# Minimal non-empty snapshot so the factory's egress PEP can run. The
# factory now requires a snapshot (no_snapshot fails closed); under the
# default ``both`` scope a local retriever is permitted, so this mirrors a
# real caller without pulling in the full settings infrastructure.
_BOTH_SNAPSHOT = {"policy.egress_scope": {"value": "both"}}


class TestRetriever(BaseRetriever):
    """Test retriever."""

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

        arbitrary_types_allowed = True

    def _get_relevant_documents(self, query: str):
        return [Document(page_content=f"Result for {query}")]

    async def _aget_relevant_documents(self, query: str):
        return self._get_relevant_documents(query)


class TestFactoryIntegration:
    """Test search engine factory integration with retrievers."""

    def setup_method(self):
        """Clear registry before each test."""
        retriever_registry.clear()

    def teardown_method(self):
        """Clear registry after each test."""
        retriever_registry.clear()

    def test_create_retriever_search_engine(self):
        """Test creating a search engine from registered retriever."""
        # Register a retriever
        retriever = TestRetriever()
        retriever_registry.register("test_retriever", retriever)

        # Create search engine
        engine = create_search_engine(
            "test_retriever", settings_snapshot=_BOTH_SNAPSHOT
        )

        assert engine is not None
        assert isinstance(engine, RetrieverSearchEngine)
        assert engine.name == "test_retriever"

        # Test it works
        results = engine.run("test query")
        assert len(results) == 1
        assert "Result for test query" in results[0]["snippet"]

    def test_retriever_not_found_fallback(self):
        """Test fallback when retriever not found."""
        # Don't register any retriever
        # The factory will require settings_snapshot when retriever not found
        with pytest.raises(RuntimeError, match="settings_snapshot is required"):
            create_search_engine("nonexistent_retriever")

    def test_create_with_max_results(self):
        """Test creating retriever engine with max_results."""
        retriever = TestRetriever()
        retriever_registry.register("test", retriever)

        engine = create_search_engine(
            "test", max_results=5, settings_snapshot=_BOTH_SNAPSHOT
        )

        assert isinstance(engine, RetrieverSearchEngine)
        assert engine.max_results == 5

    def test_multiple_retrievers_in_factory(self):
        """Test factory can handle multiple registered retrievers."""
        # Register multiple
        for i in range(3):
            retriever_registry.register(f"retriever_{i}", TestRetriever())

        # Create engines for each
        engines = []
        for i in range(3):
            engine = create_search_engine(
                f"retriever_{i}", settings_snapshot=_BOTH_SNAPSHOT
            )
            engines.append(engine)

        assert all(isinstance(e, RetrieverSearchEngine) for e in engines)
        assert engines[0].name == "retriever_0"
        assert engines[1].name == "retriever_1"
        assert engines[2].name == "retriever_2"

    def test_retrievers_in_search_config(self):
        """Test that registered retrievers appear in search config."""
        # Register retrievers
        retriever_registry.register("custom_kb", TestRetriever())
        retriever_registry.register("vector_db", TestRetriever())

        # Import after registration to trigger config update
        from local_deep_research.web_search_engines.search_engines_config import (
            search_config,
        )

        # Get fresh config
        config = search_config()

        # Check retrievers are in config
        assert "custom_kb" in config
        assert "vector_db" in config

        # Check config structure
        assert config["custom_kb"]["class_name"] == "RetrieverSearchEngine"
        assert config["custom_kb"]["is_retriever"]
        assert not config["custom_kb"]["requires_api_key"]
        assert not config["custom_kb"]["requires_llm"]