File size: 1,254 Bytes
64d7fdf
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
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)}")