diff --git a/tests/unit/test_query/test_doc_embeddings_milvus_query.py b/tests/unit/test_query/test_doc_embeddings_milvus_query.py
index ae615bff..622529e5 100644
--- a/tests/unit/test_query/test_doc_embeddings_milvus_query.py
+++ b/tests/unit/test_query/test_doc_embeddings_milvus_query.py
@@ -118,10 +118,10 @@ class TestMilvusDocEmbeddingsQueryProcessor:
         
         result = await processor.query_document_embeddings(query)
         
-        # Verify search was called twice with correct parameters
+        # Verify search was called twice with correct parameters including user/collection
         expected_calls = [
-            (([0.1, 0.2, 0.3],), {"limit": 3}),
-            (([0.4, 0.5, 0.6],), {"limit": 3}),
+            (([0.1, 0.2, 0.3], 'test_user', 'test_collection'), {"limit": 3}),
+            (([0.4, 0.5, 0.6], 'test_user', 'test_collection'), {"limit": 3}),
         ]
         assert processor.vecstore.search.call_count == 2
         for i, (expected_args, expected_kwargs) in enumerate(expected_calls):
@@ -157,7 +157,9 @@ class TestMilvusDocEmbeddingsQueryProcessor:
         result = await processor.query_document_embeddings(query)
         
         # Verify search was called with the specified limit
-        processor.vecstore.search.assert_called_once_with([0.1, 0.2, 0.3], limit=2)
+        processor.vecstore.search.assert_called_once_with(
+            [0.1, 0.2, 0.3], 'test_user', 'test_collection', limit=2
+        )
         
         # Verify all results are returned (Milvus handles limit internally)
         assert len(result) == 4
@@ -196,7 +198,9 @@ class TestMilvusDocEmbeddingsQueryProcessor:
         result = await processor.query_document_embeddings(query)
         
         # Verify search was called
-        processor.vecstore.search.assert_called_once_with([0.1, 0.2, 0.3], limit=5)
+        processor.vecstore.search.assert_called_once_with(
+            [0.1, 0.2, 0.3], 'test_user', 'test_collection', limit=5
+        )
         
         # Verify empty results
         assert len(result) == 0
diff --git a/tests/unit/test_query/test_graph_embeddings_milvus_query.py b/tests/unit/test_query/test_graph_embeddings_milvus_query.py
index 5fbb74d5..ebacfaaf 100644
--- a/tests/unit/test_query/test_graph_embeddings_milvus_query.py
+++ b/tests/unit/test_query/test_graph_embeddings_milvus_query.py
@@ -133,8 +133,10 @@ class TestMilvusGraphEmbeddingsQueryProcessor:
         
         result = await processor.query_graph_embeddings(query)
         
-        # Verify search was called with correct parameters
-        processor.vecstore.search.assert_called_once_with([0.1, 0.2, 0.3], limit=10)
+        # Verify search was called with correct parameters including user/collection
+        processor.vecstore.search.assert_called_once_with(
+            [0.1, 0.2, 0.3], 'test_user', 'test_collection', limit=10
+        )
         
         # Verify results are converted to Value objects
         assert len(result) == 3
@@ -171,10 +173,10 @@ class TestMilvusGraphEmbeddingsQueryProcessor:
         
         result = await processor.query_graph_embeddings(query)
         
-        # Verify search was called twice with correct parameters
+        # Verify search was called twice with correct parameters including user/collection
         expected_calls = [
-            (([0.1, 0.2, 0.3],), {"limit": 6}),
-            (([0.4, 0.5, 0.6],), {"limit": 6}),
+            (([0.1, 0.2, 0.3], 'test_user', 'test_collection'), {"limit": 6}),
+            (([0.4, 0.5, 0.6], 'test_user', 'test_collection'), {"limit": 6}),
         ]
         assert processor.vecstore.search.call_count == 2
         for i, (expected_args, expected_kwargs) in enumerate(expected_calls):
@@ -211,7 +213,9 @@ class TestMilvusGraphEmbeddingsQueryProcessor:
         result = await processor.query_graph_embeddings(query)
         
         # Verify search was called with 2*limit for better deduplication
-        processor.vecstore.search.assert_called_once_with([0.1, 0.2, 0.3], limit=4)
+        processor.vecstore.search.assert_called_once_with(
+            [0.1, 0.2, 0.3], 'test_user', 'test_collection', limit=4
+        )
         
         # Verify results are limited to the requested limit
         assert len(result) == 2
@@ -269,7 +273,9 @@ class TestMilvusGraphEmbeddingsQueryProcessor:
         result = await processor.query_graph_embeddings(query)
         
         # Verify only first vector was searched (limit reached)
-        processor.vecstore.search.assert_called_once_with([0.1, 0.2, 0.3], limit=4)
+        processor.vecstore.search.assert_called_once_with(
+            [0.1, 0.2, 0.3], 'test_user', 'test_collection', limit=4
+        )
         
         # Verify results are limited
         assert len(result) == 2
@@ -308,7 +314,9 @@ class TestMilvusGraphEmbeddingsQueryProcessor:
         result = await processor.query_graph_embeddings(query)
         
         # Verify search was called
-        processor.vecstore.search.assert_called_once_with([0.1, 0.2, 0.3], limit=10)
+        processor.vecstore.search.assert_called_once_with(
+            [0.1, 0.2, 0.3], 'test_user', 'test_collection', limit=10
+        )
         
         # Verify empty results
         assert len(result) == 0