using milvus lite to avoid connection failing issues

Author: jinhonglin-ryan

tests/vector_db/test_milvus.py

@@ -1,6 +1,10 @@
 import unittest
 from unittest.mock import patch, MagicMock
 import numpy as np
+import warnings
+
+# Filter out the pkg_resources deprecation warning from milvus_lite
+warnings.filterwarnings("ignore", category=DeprecationWarning, module="pkg_resources")
 
 from deepsearcher.vector_db import Milvus
 from deepsearcher.loader.splitter import Chunk
@@ -10,20 +14,11 @@
 class TestMilvus(unittest.TestCase):
     """Simple tests for the Milvus vector database implementation."""
 
-    @patch('pymilvus.connections')
-    @patch('pymilvus.MilvusClient')
-    def test_init(self, mock_client_class, mock_connections):
+    def test_init(self):
         """Test basic initialization."""
-        mock_client = MagicMock()
-        mock_client_class.return_value = mock_client
-        # Mock the connections.connect to prevent actual connection attempts
-        mock_connections.connect = MagicMock()
-        
         milvus = Milvus(
             default_collection="test_collection",
-            uri="http://localhost:19530",
-            token="root:Milvus",
-            db="default",
+            uri="./milvus.db",
             hybrid=False
         )
 
@@ -32,23 +27,9 @@ def test_init(self, mock_client_class, mock_connections):
         self.assertFalse(milvus.hybrid)
         self.assertIsNotNone(milvus.client)
 
-    @patch('pymilvus.connections')
-    @patch('pymilvus.MilvusClient')
-    def test_init_collection(self, mock_client_class, mock_connections):
+    def test_init_collection(self):
         """Test collection initialization."""
-        mock_client = MagicMock()
-        mock_client_class.return_value = mock_client
-        # Mock the connections.connect to prevent actual connection attempts
-        mock_connections.connect = MagicMock()
-        mock_client.has_collection.return_value = False  # Collection doesn't exist
-        
-        # Mock schema and index creation
-        mock_schema = MagicMock()
-        mock_index_params = MagicMock()
-        mock_client.create_schema.return_value = mock_schema
-        mock_client.prepare_index_params.return_value = mock_index_params
-        
-        milvus = Milvus()
+        milvus = Milvus(uri="./milvus.db")
 
         # Test collection initialization
         d = 8
@@ -63,17 +44,9 @@ def test_init_collection(self, mock_client_class, mock_connections):
 
         self.assertTrue(test_passed, "init_collection should work")
 
-    @patch('pymilvus.connections')
-    @patch('pymilvus.MilvusClient')
-    def test_insert_data_with_retrieval_results(self, mock_client_class, mock_connections):
+    def test_insert_data_with_retrieval_results(self):
         """Test inserting data using RetrievalResult objects."""
-        mock_client = MagicMock()
-        mock_client_class.return_value = mock_client
-        # Mock the connections.connect to prevent actual connection attempts
-        mock_connections.connect = MagicMock()
-        mock_client.insert.return_value = None
-        
-        milvus = Milvus()
+        milvus = Milvus(uri="./milvus.db")
 
         # Create test data
         d = 8
@@ -105,44 +78,14 @@ def test_insert_data_with_retrieval_results(self, mock_client_class, mock_connec
 
         self.assertTrue(test_passed, "insert_data should work with RetrievalResult objects")
 
-    @patch('pymilvus.connections')
-    @patch('pymilvus.MilvusClient')
-    def test_search_data(self, mock_client_class, mock_connections):
+    def test_search_data(self):
         """Test search functionality."""
-        mock_client = MagicMock()
-        mock_client_class.return_value = mock_client
-        # Mock the connections.connect to prevent actual connection attempts
-        mock_connections.connect = MagicMock()
-        
-        # Mock search results
-        d = 8
-        rng = np.random.default_rng(seed=19530)
-        mock_search_results = [[
-            {
-                "entity": {
-                    "embedding": rng.random((1, d))[0].tolist(),
-                    "text": "hello world",
-                    "reference": "local file: hi.txt",
-                    "metadata": {"a": 1}
-                },
-                "distance": 0.5
-            },
-            {
-                "entity": {
-                    "embedding": rng.random((1, d))[0].tolist(),
-                    "text": "hello milvus",
-                    "reference": "local file: hi.txt",
-                    "metadata": {"a": 1}
-                },
-                "distance": 0.8
-            }
-        ]]
-        mock_client.search.return_value = mock_search_results
-        
-        milvus = Milvus()
+        milvus = Milvus(uri="./milvus.db")
 
         # Test search
+        d = 8
         collection = "hello_deepsearcher"
+        rng = np.random.default_rng(seed=19530)
         query_vector = rng.random((1, d))[0]
 
         try:
@@ -159,22 +102,12 @@ def test_search_data(self, mock_client_class, mock_connections):
         self.assertTrue(test_passed, "search_data should work")
         if test_passed:
             self.assertIsInstance(top_2, list)
-            self.assertEqual(len(top_2), 2)
-            # Verify results are RetrievalResult objects
-            for result in top_2:
-                self.assertIsInstance(result, RetrievalResult)
+            # Note: In an empty collection, we might not get 2 results
+            self.assertIsInstance(top_2[0], RetrievalResult) if top_2 else None
 
-    @patch('pymilvus.connections')
-    @patch('pymilvus.MilvusClient')
-    def test_clear_collection(self, mock_client_class, mock_connections):
+    def test_clear_collection(self):
         """Test clearing collection."""
-        mock_client = MagicMock()
-        mock_client_class.return_value = mock_client
-        # Mock the connections.connect to prevent actual connection attempts
-        mock_connections.connect = MagicMock()
-        mock_client.drop_collection.return_value = None
-        
-        milvus = Milvus()
+        milvus = Milvus(uri="./milvus.db")
 
         collection = "hello_deepsearcher"
 
@@ -187,25 +120,9 @@ def test_clear_collection(self, mock_client_class, mock_connections):
 
         self.assertTrue(test_passed, "clear_db should work")
 
-    @patch('pymilvus.connections')
-    @patch('pymilvus.MilvusClient')
-    def test_list_collections(self, mock_client_class, mock_connections):
+    def test_list_collections(self):
         """Test listing collections."""
-        mock_client = MagicMock()
-        mock_client_class.return_value = mock_client
-        # Mock the connections.connect to prevent actual connection attempts
-        mock_connections.connect = MagicMock()
-        mock_client.list_collections.return_value = ["hello_deepsearcher", "test_collection"]
-        
-        # Mock describe_collection for each collection
-        def mock_describe(collection_name):
-            return {
-                "description": f"Description for {collection_name}",
-                "fields": []
-            }
-        mock_client.describe_collection.side_effect = mock_describe
-        
-        milvus = Milvus()
+        milvus = Milvus(uri="./milvus.db")
 
         try:
             collections = milvus.list_collections()