tests/test_retriever.py

import pytest
from langchain_core.documents import Document
from app.core.retriever import HybridRetriever


@pytest.mark.unit
class TestHybridRetriever:
    
    def test_retriever_initialization(self):
        retriever = HybridRetriever()
        
        assert retriever is not None
        assert retriever.k == 10
        assert retriever._initialized == False
    
    def test_vector_store_lazy_initialization(self):
        retriever = HybridRetriever()
        
        assert retriever._initialized == False
        
        retriever._initialize_vector_store()
        
        assert retriever._initialized == True
        assert retriever.vector_store is not None
    
    @pytest.mark.asyncio
    @pytest.mark.slow
    async def test_retrieve_documents(self):
        retriever = HybridRetriever()
        retriever._initialize_vector_store()
        
        query = "What is attention mechanism?"
        
        try:
            documents = await retriever.ainvoke(query)
            
            assert isinstance(documents, list)
            if len(documents) > 0:
                assert all(isinstance(doc, Document) for doc in documents)
        except Exception as e:
            pytest.skip(f"Retrieval test skipped: {str(e)}")