tests/test_retrieval_strategies.py

"""Tests for retrieval strategies."""

import tempfile
from unittest.mock import MagicMock, patch

import pytest
from langchain_core.documents import Document

from src.retrieval.strategies.bm25 import BM25Strategy
from src.retrieval.strategies.bm25_vector import BM25VectorStrategy
from src.retrieval.strategies.vector import VectorStrategy


@pytest.fixture
def sample_documents():
    """Create sample documents for testing."""
    return [
        Document(
            page_content="Tin Hoang is an AI Research Engineer with expertise in machine learning.",
            metadata={"source": "profile.md"},
        ),
        Document(
            page_content="He has experience with Python, PyTorch, and TensorFlow frameworks.",
            metadata={"source": "skills.md"},
        ),
        Document(
            page_content="His research focuses on federated learning for medical AI applications.",
            metadata={"source": "research.md"},
        ),
        Document(
            page_content="He developed OCR systems using deep learning techniques.",
            metadata={"source": "projects.md"},
        ),
    ]


@pytest.fixture
def temp_index_path(tmp_path):
    """Create temporary paths for indexes."""
    return {
        "bm25": str(tmp_path / "bm25_test"),
        "vector": str(tmp_path / "vector_test"),
    }


@pytest.fixture
def mock_vectorstore_manager():
    """Create a mock VectorStoreManager."""
    # VectorStoreManager is imported inside VectorStrategy using lazy import
    # from src.vectorstore, so we patch it there
    with patch("src.vectorstore.VectorStoreManager") as mock:
        mock_manager = MagicMock()
        mock_manager.persist_directory = tempfile.mkdtemp()
        mock.return_value = mock_manager
        yield mock_manager


# --- BM25Strategy Tests ---


def test_bm25_init(temp_index_path):
    """Test BM25Strategy initialization."""
    config = {
        "retrieval": {
            "bm25": {
                "k": 5,
                "persist_path": temp_index_path["bm25"],
                "tokenizer": "simple",
            }
        }
    }
    strategy = BM25Strategy(config)

    assert strategy.name == "bm25"
    assert strategy.k == 5
    assert not strategy.is_index_ready()


def test_bm25_build_and_retrieve(sample_documents, temp_index_path):
    """Test building index and retrieving documents."""
    config = {
        "retrieval": {
            "bm25": {
                "k": 3,
                "persist_path": temp_index_path["bm25"],
                "tokenizer": "simple",
            }
        }
    }
    strategy = BM25Strategy(config)

    # Build index
    strategy.build_index(sample_documents)
    assert strategy.is_index_ready()

    # Retrieve
    results = strategy.retrieve("machine learning AI", k=2)
    assert len(results) <= 2
    assert all(isinstance(doc, Document) for doc in results)


def test_bm25_load_index(sample_documents, temp_index_path):
    """Test loading a saved index."""
    config = {
        "retrieval": {
            "bm25": {
                "k": 3,
                "persist_path": temp_index_path["bm25"],
            }
        }
    }

    # Build and save
    strategy1 = BM25Strategy(config)
    strategy1.build_index(sample_documents)

    # Load in new instance
    strategy2 = BM25Strategy(config)
    assert strategy2.load_index()
    assert strategy2.is_index_ready()

    # Should be able to retrieve
    results = strategy2.retrieve("Python", k=2)
    assert len(results) > 0


def test_bm25_as_retriever(sample_documents, temp_index_path):
    """Test getting LangChain retriever."""
    config = {
        "retrieval": {
            "bm25": {
                "k": 2,
                "persist_path": temp_index_path["bm25"],
            }
        }
    }
    strategy = BM25Strategy(config)
    strategy.build_index(sample_documents)

    retriever = strategy.as_retriever()
    assert retriever is not None

    # Should work with invoke
    results = retriever.invoke("federated learning")
    assert isinstance(results, list)


def test_bm25_get_index_stats(sample_documents, temp_index_path):
    """Test getting index statistics."""
    config = {
        "retrieval": {
            "bm25": {
                "k": 3,
                "persist_path": temp_index_path["bm25"],
            }
        }
    }
    strategy = BM25Strategy(config)
    strategy.build_index(sample_documents)

    stats = strategy.get_index_stats()
    assert stats["strategy"] == "bm25"
    assert stats["num_documents"] == 4
    assert stats["is_built"]


# --- BM25VectorStrategy Tests ---


def test_bm25_vector_init(temp_index_path):
    """Test BM25VectorStrategy initialization."""
    config = {
        "retrieval": {
            "final_k": 4,
            "vector": {"search_type": "similarity", "k": 10},
            "bm25": {"k": 10, "persist_path": temp_index_path["bm25"]},
            "fusion": {
                "algorithm": "rrf",
                "rrf_k": 60,
                "weights": {"vector": 0.7, "bm25": 0.3},
            },
        }
    }
    strategy = BM25VectorStrategy(config)

    assert strategy.name == "bm25_vector"
    assert strategy.vector_weight == 0.7
    assert strategy.bm25_weight == 0.3
    assert strategy.final_k == 4


@patch("src.retrieval.strategies.vector.VectorStrategy.build_index")
@patch("src.retrieval.strategies.vector.VectorStrategy.load_index")
def test_bm25_vector_build_index_calls_both(
    mock_vector_load, mock_vector_build, sample_documents, temp_index_path
):
    """Test that build_index builds both BM25 and Vector indexes."""
    mock_vector_load.return_value = True

    config = {
        "retrieval": {
            "final_k": 4,
            "vector": {"search_type": "similarity", "k": 10},
            "bm25": {"k": 10, "persist_path": temp_index_path["bm25"]},
            "fusion": {
                "algorithm": "rrf",
                "rrf_k": 60,
                "weights": {"vector": 0.7, "bm25": 0.3},
            },
        }
    }
    strategy = BM25VectorStrategy(config)
    strategy.build_index(sample_documents)

    # Vector build should be called
    mock_vector_build.assert_called_once_with(sample_documents)

    # BM25 should be built (we can check the underlying store)
    assert strategy.bm25_strategy.is_index_ready()


def test_bm25_vector_get_index_stats(temp_index_path):
    """Test getting combined index statistics."""
    config = {
        "retrieval": {
            "final_k": 4,
            "vector": {"search_type": "similarity", "k": 10},
            "bm25": {"k": 10, "persist_path": temp_index_path["bm25"]},
            "fusion": {
                "algorithm": "rrf",
                "rrf_k": 60,
                "weights": {"vector": 0.7, "bm25": 0.3},
            },
        }
    }
    strategy = BM25VectorStrategy(config)

    stats = strategy.get_index_stats()
    assert stats["strategy"] == "bm25_vector"
    assert "vector" in stats
    assert "bm25" in stats
    assert "fusion" in stats
    assert stats["fusion"]["algorithm"] == "rrf"
    assert stats["fusion"]["weights"]["vector"] == 0.7


@patch("src.retrieval.strategies.vector.VectorStrategy.load_index")
@patch("src.retrieval.strategies.vector.VectorStrategy.is_index_ready")
@patch("src.retrieval.strategies.vector.VectorStrategy.retrieve")
def test_bm25_vector_retrieve_with_fusion(
    mock_vector_retrieve,
    mock_vector_ready,
    mock_vector_load,
    sample_documents,
    temp_index_path,
):
    """Test retrieve combines vector and BM25 results."""
    mock_vector_load.return_value = True
    mock_vector_ready.return_value = True
    mock_vector_retrieve.return_value = sample_documents[:2]

    config = {
        "retrieval": {
            "final_k": 3,
            "vector": {"search_type": "similarity", "k": 10},
            "bm25": {"k": 10, "persist_path": temp_index_path["bm25"]},
            "fusion": {
                "algorithm": "rrf",
                "rrf_k": 60,
                "weights": {"vector": 0.7, "bm25": 0.3},
            },
        }
    }
    strategy = BM25VectorStrategy(config)
    strategy.bm25_strategy.build_index(sample_documents)
    strategy._is_initialized = True

    results = strategy.retrieve("machine learning AI", k=3)

    assert len(results) <= 3
    assert all(isinstance(doc, Document) for doc in results)


@patch("src.retrieval.strategies.vector.VectorStrategy.load_index")
@patch("src.retrieval.strategies.vector.VectorStrategy.is_index_ready")
@patch("src.retrieval.strategies.vector.VectorStrategy.retrieve")
def test_bm25_vector_retrieve_vector_only_fallback(
    mock_vector_retrieve,
    mock_vector_ready,
    mock_vector_load,
    sample_documents,
    temp_index_path,
):
    """Test retrieve falls back to vector only when BM25 fails."""
    mock_vector_load.return_value = True
    mock_vector_ready.return_value = True
    mock_vector_retrieve.return_value = sample_documents[:2]

    config = {
        "retrieval": {
            "final_k": 3,
            "vector": {"search_type": "similarity", "k": 10},
            "bm25": {"k": 10, "persist_path": temp_index_path["bm25"]},
            "fusion": {
                "algorithm": "rrf",
                "rrf_k": 60,
                "weights": {"vector": 0.7, "bm25": 0.3},
            },
        }
    }
    strategy = BM25VectorStrategy(config)
    strategy._is_initialized = True
    # Don't build BM25 index - it should return empty

    results = strategy.retrieve("query", k=2)

    # Should return vector results as fallback
    assert len(results) == 2


@patch("src.retrieval.strategies.vector.VectorStrategy.load_index")
@patch("src.retrieval.strategies.vector.VectorStrategy.is_index_ready")
def test_bm25_vector_retrieve_bm25_only_fallback(
    mock_vector_ready, mock_vector_load, sample_documents, temp_index_path
):
    """Test retrieve falls back to BM25 only when vector fails."""
    mock_vector_load.return_value = True
    mock_vector_ready.return_value = False  # Vector not ready

    config = {
        "retrieval": {
            "final_k": 3,
            "vector": {"search_type": "similarity", "k": 10},
            "bm25": {"k": 10, "persist_path": temp_index_path["bm25"]},
            "fusion": {
                "algorithm": "rrf",
                "rrf_k": 60,
                "weights": {"vector": 0.7, "bm25": 0.3},
            },
        }
    }
    strategy = BM25VectorStrategy(config)
    strategy.bm25_strategy.build_index(sample_documents)
    strategy._is_initialized = True

    results = strategy.retrieve("machine learning", k=2)

    # Should return BM25 results as fallback
    assert len(results) <= 2


@patch("src.retrieval.strategies.vector.VectorStrategy.load_index")
@patch("src.retrieval.strategies.vector.VectorStrategy.is_index_ready")
def test_bm25_vector_retrieve_both_empty(mock_vector_ready, mock_vector_load, temp_index_path):
    """Test retrieve returns empty when both fail."""
    mock_vector_load.return_value = True
    mock_vector_ready.return_value = False

    config = {
        "retrieval": {
            "final_k": 3,
            "vector": {"search_type": "similarity", "k": 10},
            "bm25": {"k": 10, "persist_path": temp_index_path["bm25"]},
            "fusion": {"algorithm": "rrf", "rrf_k": 60},
        }
    }
    strategy = BM25VectorStrategy(config)
    strategy._is_initialized = True
    # Neither index is ready

    results = strategy.retrieve("query", k=2)

    assert results == []


def test_bm25_vector_retrieve_not_initialized_raises(temp_index_path):
    """Test retrieve raises when not initialized and can't load."""
    config = {
        "retrieval": {
            "final_k": 3,
            "vector": {"search_type": "similarity", "k": 10},
            "bm25": {"k": 10, "persist_path": temp_index_path["bm25"]},
            "fusion": {"algorithm": "rrf", "rrf_k": 60},
        }
    }
    strategy = BM25VectorStrategy(config)

    # Mock load_index to return False
    with patch.object(strategy, "load_index", return_value=False):
        with pytest.raises(ValueError, match="Indexes not available"):
            strategy.retrieve("query")


@patch("src.retrieval.strategies.vector.VectorStrategy.load_index")
def test_bm25_vector_load_index_both_success(mock_vector_load, sample_documents, temp_index_path):
    """Test load_index succeeds when both indexes load."""
    mock_vector_load.return_value = True

    config = {
        "retrieval": {
            "final_k": 4,
            "vector": {"search_type": "similarity", "k": 10},
            "bm25": {"k": 10, "persist_path": temp_index_path["bm25"]},
            "fusion": {"algorithm": "rrf", "rrf_k": 60},
        }
    }
    strategy = BM25VectorStrategy(config)

    # Build BM25 first so it can be loaded
    strategy.bm25_strategy.build_index(sample_documents)

    # Create new strategy instance
    strategy2 = BM25VectorStrategy(config)
    result = strategy2.load_index()

    assert result is True
    assert strategy2._is_initialized


@patch("src.retrieval.strategies.vector.VectorStrategy.load_index")
def test_bm25_vector_load_index_vector_only(mock_vector_load, temp_index_path):
    """Test load_index with only vector succeeding."""
    mock_vector_load.return_value = True

    config = {
        "retrieval": {
            "final_k": 4,
            "vector": {"search_type": "similarity", "k": 10},
            "bm25": {"k": 10, "persist_path": temp_index_path["bm25"]},
            "fusion": {"algorithm": "rrf", "rrf_k": 60},
        }
    }
    strategy = BM25VectorStrategy(config)

    result = strategy.load_index()

    # Should succeed with partial functionality
    assert result is True
    assert strategy._is_initialized


@patch("src.retrieval.strategies.vector.VectorStrategy.load_index")
def test_bm25_vector_load_index_bm25_only(mock_vector_load, sample_documents, temp_index_path):
    """Test load_index with only BM25 succeeding."""
    mock_vector_load.return_value = False

    config = {
        "retrieval": {
            "final_k": 4,
            "vector": {"search_type": "similarity", "k": 10},
            "bm25": {"k": 10, "persist_path": temp_index_path["bm25"]},
            "fusion": {"algorithm": "rrf", "rrf_k": 60},
        }
    }
    strategy = BM25VectorStrategy(config)
    strategy.bm25_strategy.build_index(sample_documents)

    strategy2 = BM25VectorStrategy(config)
    result = strategy2.load_index()

    assert result is True
    assert strategy2._is_initialized


@patch("src.retrieval.strategies.vector.VectorStrategy.load_index")
def test_bm25_vector_load_index_both_fail(mock_vector_load, temp_index_path):
    """Test load_index fails when both fail."""
    mock_vector_load.return_value = False

    config = {
        "retrieval": {
            "final_k": 4,
            "vector": {"search_type": "similarity", "k": 10},
            "bm25": {"k": 10, "persist_path": temp_index_path["bm25"]},
            "fusion": {"algorithm": "rrf", "rrf_k": 60},
        }
    }
    strategy = BM25VectorStrategy(config)

    result = strategy.load_index()

    assert result is False


@patch("src.retrieval.strategies.vector.VectorStrategy.load_index")
@patch("src.retrieval.strategies.vector.VectorStrategy.is_index_ready")
@patch("src.retrieval.strategies.vector.VectorStrategy.as_retriever")
def test_bm25_vector_as_retriever(
    mock_vector_retriever,
    mock_vector_ready,
    mock_vector_load,
    sample_documents,
    temp_index_path,
):
    """Test as_retriever creates FusionRetriever."""
    from langchain_core.retrievers import BaseRetriever

    mock_vector_load.return_value = True
    mock_vector_ready.return_value = True
    # Use spec to satisfy Pydantic validation
    mock_retriever = MagicMock(spec=BaseRetriever)
    mock_vector_retriever.return_value = mock_retriever

    config = {
        "retrieval": {
            "final_k": 4,
            "vector": {"search_type": "similarity", "k": 10},
            "bm25": {"k": 10, "persist_path": temp_index_path["bm25"]},
            "fusion": {"algorithm": "rrf", "rrf_k": 60},
        }
    }
    strategy = BM25VectorStrategy(config)
    strategy.bm25_strategy.build_index(sample_documents)
    strategy._is_initialized = True

    retriever = strategy.as_retriever()

    assert retriever is not None


@patch("src.retrieval.strategies.vector.VectorStrategy.load_index")
@patch("src.retrieval.strategies.vector.VectorStrategy.is_index_ready")
def test_bm25_vector_as_retriever_no_retrievers_raises(
    mock_vector_ready, mock_vector_load, temp_index_path
):
    """Test as_retriever raises when no retrievers available."""
    mock_vector_load.return_value = True
    mock_vector_ready.return_value = False

    config = {
        "retrieval": {
            "final_k": 4,
            "vector": {"search_type": "similarity", "k": 10},
            "bm25": {"k": 10, "persist_path": temp_index_path["bm25"]},
            "fusion": {"algorithm": "rrf", "rrf_k": 60},
        }
    }
    strategy = BM25VectorStrategy(config)
    strategy._is_initialized = True

    with pytest.raises(ValueError, match="No retrievers available"):
        strategy.as_retriever()


@patch("src.retrieval.strategies.vector.VectorStrategy.load_index")
def test_bm25_vector_delete_index(mock_vector_load, sample_documents, temp_index_path):
    """Test delete_index clears indexes."""
    mock_vector_load.return_value = True

    config = {
        "retrieval": {
            "final_k": 4,
            "vector": {"search_type": "similarity", "k": 10},
            "bm25": {"k": 10, "persist_path": temp_index_path["bm25"]},
            "fusion": {"algorithm": "rrf", "rrf_k": 60},
        }
    }
    strategy = BM25VectorStrategy(config)
    strategy.bm25_strategy.build_index(sample_documents)
    strategy._is_initialized = True

    strategy.delete_index()

    assert not strategy._is_initialized


# --- VectorStrategy Tests ---


def test_vector_init(mock_vectorstore_manager):
    """Test VectorStrategy initialization."""
    config = {
        "retrieval": {
            "vector": {
                "search_type": "similarity",
                "k": 4,
                "search_kwargs": {},
            }
        }
    }
    strategy = VectorStrategy(config)

    assert strategy.name == "vector"
    assert strategy.k == 4
    assert strategy.search_type == "similarity"


def test_vector_build_index(mock_vectorstore_manager, sample_documents):
    """Test building vector index."""
    config = {"retrieval": {"vector": {"search_type": "similarity", "k": 4}}}
    strategy = VectorStrategy(config)

    strategy.build_index(sample_documents)

    mock_vectorstore_manager.create_vectorstore.assert_called_once_with(sample_documents)
    assert strategy._is_initialized


def test_vector_load_index_success(mock_vectorstore_manager):
    """Test loading vector index successfully."""
    config = {"retrieval": {"vector": {"search_type": "similarity", "k": 4}}}

    with patch("src.retrieval.strategies.vector.Path") as mock_path:
        mock_path.return_value.exists.return_value = True

        strategy = VectorStrategy(config)
        result = strategy.load_index()

        assert result is True
        assert strategy._is_initialized
        mock_vectorstore_manager.load_vectorstore.assert_called_once()


def test_vector_load_index_not_found(mock_vectorstore_manager):
    """Test loading when index not found."""
    config = {"retrieval": {"vector": {"search_type": "similarity", "k": 4}}}

    with patch("src.retrieval.strategies.vector.Path") as mock_path:
        mock_path.return_value.exists.return_value = False

        strategy = VectorStrategy(config)
        result = strategy.load_index()

        assert result is False


def test_vector_load_index_error(mock_vectorstore_manager):
    """Test loading with error."""
    config = {"retrieval": {"vector": {"search_type": "similarity", "k": 4}}}
    mock_vectorstore_manager.load_vectorstore.side_effect = Exception("Load error")

    with patch("src.retrieval.strategies.vector.Path") as mock_path:
        mock_path.return_value.exists.return_value = True

        strategy = VectorStrategy(config)
        result = strategy.load_index()

        assert result is False


def test_vector_retrieve(mock_vectorstore_manager, sample_documents):
    """Test retrieving documents."""
    config = {"retrieval": {"vector": {"search_type": "similarity", "k": 4}}}
    mock_vectorstore_manager.similarity_search.return_value = sample_documents[:2]

    strategy = VectorStrategy(config)
    strategy._is_initialized = True

    results = strategy.retrieve("machine learning", k=2)

    assert len(results) == 2
    mock_vectorstore_manager.similarity_search.assert_called_once_with("machine learning", k=2)


def test_vector_retrieve_loads_if_not_initialized(mock_vectorstore_manager, sample_documents):
    """Test retrieve loads index if not initialized."""
    config = {"retrieval": {"vector": {"search_type": "similarity", "k": 4}}}
    mock_vectorstore_manager.similarity_search.return_value = sample_documents[:2]

    with patch("src.retrieval.strategies.vector.Path") as mock_path:
        mock_path.return_value.exists.return_value = True

        strategy = VectorStrategy(config)
        strategy._is_initialized = False

        strategy.retrieve("query", k=2)

        mock_vectorstore_manager.load_vectorstore.assert_called_once()


def test_vector_as_retriever(mock_vectorstore_manager):
    """Test getting LangChain retriever."""
    config = {"retrieval": {"vector": {"search_type": "similarity", "k": 4}}}
    mock_retriever = MagicMock()
    mock_vectorstore_manager.get_retriever.return_value = mock_retriever

    strategy = VectorStrategy(config)
    strategy._is_initialized = True

    retriever = strategy.as_retriever()

    assert retriever is mock_retriever


def test_vector_as_retriever_with_kwargs(mock_vectorstore_manager):
    """Test getting retriever with custom kwargs."""
    config = {"retrieval": {"vector": {"search_type": "similarity", "k": 4}}}
    mock_retriever = MagicMock()
    mock_vectorstore_manager.get_retriever.return_value = mock_retriever

    strategy = VectorStrategy(config)
    strategy._is_initialized = True

    strategy.as_retriever(search_type="mmr", search_kwargs={"k": 10})

    mock_vectorstore_manager.get_retriever.assert_called_once()
    call_kwargs = mock_vectorstore_manager.get_retriever.call_args.kwargs
    assert call_kwargs.get("search_type") == "mmr"


def test_vector_get_index_stats_initialized(mock_vectorstore_manager):
    """Test getting index stats when initialized."""
    config = {"retrieval": {"vector": {"search_type": "similarity", "k": 4}}}
    mock_store = MagicMock()
    mock_store._collection.count.return_value = 100
    mock_vectorstore_manager.vectorstore = mock_store
    mock_vectorstore_manager.persist_directory = "/test/path"
    mock_vectorstore_manager.collection_name = "test_collection"

    strategy = VectorStrategy(config)
    strategy._is_initialized = True

    stats = strategy.get_index_stats()

    assert stats["strategy"] == "vector"
    assert stats["num_documents"] == 100
    assert stats["collection_name"] == "test_collection"


def test_vector_get_index_stats_error(mock_vectorstore_manager):
    """Test getting index stats handles errors."""
    config = {"retrieval": {"vector": {"search_type": "similarity", "k": 4}}}
    mock_store = MagicMock()
    mock_store._collection.count.side_effect = Exception("Stats error")
    mock_vectorstore_manager.vectorstore = mock_store

    strategy = VectorStrategy(config)
    strategy._is_initialized = True

    stats = strategy.get_index_stats()

    # Should still return base stats
    assert stats["strategy"] == "vector"