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

"""Pytest configuration for retriever integration tests."""

from unittest.mock import Mock

import pytest
from langchain_core.retrievers import BaseRetriever, Document


@pytest.fixture
def mock_retriever():
    """Create a mock retriever for testing."""

    class MockRetriever(BaseRetriever):
        def __init__(self):
            self.queries = []
            self.documents = [
                Document(
                    page_content="Test document content",
                    metadata={"title": "Test Doc", "source": "test"},
                )
            ]

        def get_relevant_documents(self, query: str):
            self.queries.append(query)
            return self.documents

        async def aget_relevant_documents(self, query: str):
            return self.get_relevant_documents(query)

    return MockRetriever()


@pytest.fixture
def multiple_retrievers():
    """Create multiple mock retrievers."""
    retrievers = {}
    for i in range(3):

        class NamedRetriever(BaseRetriever):
            def __init__(self, name):
                self.name = name
                self.queries = []

            def get_relevant_documents(self, query: str):
                self.queries.append(query)
                return [
                    Document(
                        page_content=f"Content from {self.name}",
                        metadata={
                            "title": f"{self.name} Doc",
                            "source": self.name,
                        },
                    )
                ]

            async def aget_relevant_documents(self, query: str):
                return self.get_relevant_documents(query)

        retrievers[f"retriever_{i}"] = NamedRetriever(f"retriever_{i}")

    return retrievers


@pytest.fixture(autouse=True)
def cleanup_registry():
    """Clean up retriever registry before and after each test."""
    from local_deep_research.web_search_engines.retriever_registry import (
        retriever_registry,
    )

    # Clear before test
    retriever_registry.clear()

    yield

    # Clear after test
    retriever_registry.clear()


@pytest.fixture
def mock_search_system():
    """Mock the search system for API tests."""
    mock = Mock()
    mock.analyze_topic.return_value = {
        "current_knowledge": "Test summary",
        "findings": ["Finding 1", "Finding 2"],
        "iterations": 1,
        "questions": {},
        "formatted_findings": "Formatted findings",
        "all_links_of_system": ["source1", "source2"],
    }
    return mock