norasalem88/rag_assistant_devops

1import os
2import sys
3import unittest
4
5# Add project root to sys.path so we can import src modules
6BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
7sys.path.insert(0, BASE_DIR)
8
9from src.retrieval import get_retriever, retrieve_relevant_context
10
11class TestRetrieval(unittest.TestCase):
12    
13    @classmethod
14    def setUpClass(cls):
15        """Load the vectorstore once for all tests."""
16        cls.vectorstore = get_retriever()
17        
18    def test_vectorstore_loaded(self):
19        """Ensure vector store is initialized properly."""
20        self.assertIsNotNone(self.vectorstore, "Vector store should not be None; ensure ingestion has run.")
21
22    def test_single_active_document_routing(self):
23        """Ensure chunks retrieved belong ONLY to the active document."""
24        if not self.vectorstore:
25            self.skipTest("Vectorstore not loaded")
26            
27        test_query = "What is a Kubernetes Pod?"
28        
29        # We assume there are some PDFs ingested. Let's find an available doc_id or fake one.
30        # If no documents are heavily populated, the distance threshold might block all results.
31        # For this test, we just want to verify the filter works.
32        mock_active_doc_id = "test_document.pdf"
33        
34        # Execute retrieve
35        docs = retrieve_relevant_context(
36            self.vectorstore, 
37            test_query, 
38            active_doc_ids=[mock_active_doc_id], 
39            k=5, 
40            distance_threshold=2.0 # lenient threshold to just check metadata filtering
41        )
42        
43        # Verify metadata
44        for doc in docs:
45            self.assertEqual(
46                doc.metadata.get("doc_id"), 
47                mock_active_doc_id, 
48                f"Retrieved chunk from {doc.metadata.get('doc_id')}, expected {mock_active_doc_id}"
49            )
50            
51    def test_multi_active_document_routing(self):
52        """Ensure chunks retrieved belong to one of the active documents."""
53        if not self.vectorstore:
54            self.skipTest("Vectorstore not loaded")
55            
56        test_query = "Docker vs Kubernetes?"
57        active_docs = ["docker_guide.pdf", "k8s_guide.pdf"]
58        
59        docs = retrieve_relevant_context(
60            self.vectorstore, 
61            test_query, 
62            active_doc_ids=active_docs, 
63            k=5, 
64            distance_threshold=2.0
65        )
66        
67        for doc in docs:
68            self.assertIn(
69                doc.metadata.get("doc_id"), 
70                active_docs, 
71                "Retrieved chunk from a document not in the active doc list."
72            )
73            
74if __name__ == "__main__":
75    unittest.main()
76