Comprehensive test suite for FAISS functionality.
| 37 | return args |
| 38 | |
| 39 | class FaissTester: |
| 40 | """Comprehensive test suite for FAISS functionality.""" |
| 41 | |
| 42 | def __init__(self): |
| 43 | self.test_dir = None |
| 44 | self.faiss_service = None |
| 45 | self.passed_tests = 0 |
| 46 | self.failed_tests = 0 |
| 47 | self.test_document_ids = [] |
| 48 | |
| 49 | async def setup(self): |
| 50 | """Setup test environment.""" |
| 51 | self.test_dir = Path(tempfile.mkdtemp()) |
| 52 | embedding_function = model_manager.get("text-embedding-3-large") |
| 53 | self.faiss_service = FaissService( |
| 54 | base_dir=self.test_dir, |
| 55 | embedding_function=embedding_function |
| 56 | ) |
| 57 | print(f"Test directory: {self.test_dir}") |
| 58 | |
| 59 | async def cleanup(self): |
| 60 | """Cleanup test environment.""" |
| 61 | import shutil |
| 62 | if self.test_dir and self.test_dir.exists(): |
| 63 | shutil.rmtree(self.test_dir) |
| 64 | |
| 65 | def assert_test(self, condition: bool, test_name: str, message: str = ""): |
| 66 | """Assert test condition and track results.""" |
| 67 | if condition: |
| 68 | print(f"✓ {test_name}") |
| 69 | self.passed_tests += 1 |
| 70 | else: |
| 71 | print(f"✗ {test_name}: {message}") |
| 72 | self.failed_tests += 1 |
| 73 | |
| 74 | async def test_basic_document_operations(self): |
| 75 | """Test basic document add/delete operations.""" |
| 76 | print("\n=== Testing Basic Document Operations ===") |
| 77 | |
| 78 | # Test add single document |
| 79 | texts = ["This is a test document about machine learning and AI."] |
| 80 | metadatas = [{"category": "AI", "author": "test_user"}] |
| 81 | |
| 82 | request = FaissAddRequest(texts=texts, metadatas=metadatas) |
| 83 | result = await self.faiss_service.add_documents(request) |
| 84 | |
| 85 | self.assert_test(result.count == 1, |
| 86 | "Add single document", |
| 87 | f"Expected count=1, got count={result.count}") |
| 88 | |
| 89 | if result.ids: |
| 90 | self.test_document_ids.extend(result.ids) |
| 91 | |
| 92 | # Test add multiple documents |
| 93 | texts = [ |
| 94 | "Python is a great programming language for data science.", |
| 95 | "Machine learning algorithms can learn from data patterns.", |
| 96 | "Deep learning uses neural networks with multiple layers." |