clean test for vector store

Author: Gaudy Blanco

tests/integration/vector_stores/test_azure_ai_search.py

@@ -8,6 +8,7 @@
 
 import pytest
 
+from graphrag.config.models.vector_store_schema_config import VectorStoreSchemaConfig
 from graphrag.vector_stores.azure_ai_search import AzureAISearchVectorStore
 from graphrag.vector_stores.base import VectorStoreDocument
 
@@ -39,7 +40,10 @@ def mock_index_client(self):
     @pytest.fixture
     def vector_store(self, mock_search_client, mock_index_client):
         """Create an Azure AI Search vector store instance."""
-        vector_store = AzureAISearchVectorStore(collection_name="test_vectors")
+        vector_store = AzureAISearchVectorStore(
+            collection_name="test_vectors",
+            vector_store_schema_config=VectorStoreSchemaConfig(),
+        )
 
         # Create the necessary mocks first
         vector_store.db_connection = mock_search_client

tests/integration/vector_stores/test_cosmosdb.py

@@ -8,6 +8,7 @@
 import numpy as np
 import pytest
 
+from graphrag.config.models.vector_store_schema_config import VectorStoreSchemaConfig
 from graphrag.vector_stores.base import VectorStoreDocument
 from graphrag.vector_stores.cosmosdb import CosmosDBVectorStore
 
@@ -25,6 +26,7 @@ def test_vector_store_operations():
     """Test basic vector store operations with CosmosDB."""
     vector_store = CosmosDBVectorStore(
         collection_name="testvector",
+        vector_store_schema_config=VectorStoreSchemaConfig(),
     )
 
     try:
@@ -79,6 +81,7 @@ def test_clear():
     """Test clearing the vector store."""
     vector_store = CosmosDBVectorStore(
         collection_name="testclear",
+        vector_store_schema_config=VectorStoreSchemaConfig(),
     )
     try:
         vector_store.connect(

tests/integration/vector_stores/test_factory.py

@@ -8,6 +8,7 @@
 import pytest
 
 from graphrag.config.enums import VectorStoreType
+from graphrag.config.models.vector_store_schema_config import VectorStoreSchemaConfig
 from graphrag.vector_stores.azure_ai_search import AzureAISearchVectorStore
 from graphrag.vector_stores.base import BaseVectorStore
 from graphrag.vector_stores.cosmosdb import CosmosDBVectorStore
@@ -21,10 +22,12 @@ def test_create_lancedb_vector_store():
         "db_uri": "/tmp/lancedb",
     }
     vector_store = VectorStoreFactory.create_vector_store(
-        VectorStoreType.LanceDB.value, kwargs
+        vector_store_type=VectorStoreType.LanceDB.value,
+        vector_store_schema_config=VectorStoreSchemaConfig(),
+        kwargs=kwargs,
     )
     assert isinstance(vector_store, LanceDBVectorStore)
-    assert vector_store.collection_name == "test_collection"
+    assert vector_store.index_name == "test_collection"
 
 
 @pytest.mark.skip(reason="Azure AI Search requires credentials and setup")
@@ -35,7 +38,9 @@ def test_create_azure_ai_search_vector_store():
         "api_key": "test_key",
     }
     vector_store = VectorStoreFactory.create_vector_store(
-        VectorStoreType.AzureAISearch.value, kwargs
+        vector_store_type=VectorStoreType.AzureAISearch.value,
+        vector_store_schema_config=VectorStoreSchemaConfig(),
+        kwargs=kwargs,
     )
     assert isinstance(vector_store, AzureAISearchVectorStore)
 
@@ -47,9 +52,13 @@ def test_create_cosmosdb_vector_store():
         "connection_string": "AccountEndpoint=https://test.documents.azure.com:443/;AccountKey=test_key==",
         "database_name": "test_db",
     }
+
     vector_store = VectorStoreFactory.create_vector_store(
-        VectorStoreType.CosmosDB.value, kwargs
+        vector_store_type=VectorStoreType.CosmosDB.value,
+        vector_store_schema_config=VectorStoreSchemaConfig(),
+        kwargs=kwargs,
     )
+
     assert isinstance(vector_store, CosmosDBVectorStore)
 
 
@@ -67,7 +76,12 @@ def test_register_and_create_custom_vector_store():
     VectorStoreFactory.register(
         "custom", lambda **kwargs: custom_vector_store_class(**kwargs)
     )
-    vector_store = VectorStoreFactory.create_vector_store("custom", {})
+
+    vector_store = VectorStoreFactory.create_vector_store(
+        vector_store_type="custom",
+        vector_store_schema_config=VectorStoreSchemaConfig(),
+        kwargs={},
+    )
 
     assert custom_vector_store_class.called
     assert vector_store is instance
@@ -89,7 +103,11 @@ def test_get_vector_store_types():
 
 def test_create_unknown_vector_store():
     with pytest.raises(ValueError, match="Unknown vector store type: unknown"):
-        VectorStoreFactory.create_vector_store("unknown", {})
+        VectorStoreFactory.create_vector_store(
+            vector_store_type="unknown",
+            vector_store_schema_config=VectorStoreSchemaConfig(),
+            kwargs={},
+        )
 
 
 def test_is_supported_type():
@@ -139,6 +157,9 @@ def search_by_id(self, id):
 
     # Test creating an instance
     vector_store = VectorStoreFactory.create_vector_store(
-        "custom_class", {"collection_name": "test"}
+        vector_store_type="custom_class",
+        vector_store_schema_config=VectorStoreSchemaConfig(),
+        kwargs={"collection_name": "test"},
     )
+
     assert isinstance(vector_store, CustomVectorStore)

tests/integration/vector_stores/test_lancedb.py

@@ -8,6 +8,7 @@
 
 import numpy as np
 
+from graphrag.config.models.vector_store_schema_config import VectorStoreSchemaConfig
 from graphrag.vector_stores.base import VectorStoreDocument
 from graphrag.vector_stores.lancedb import LanceDBVectorStore
 
@@ -17,7 +18,11 @@ def test_vector_store_operations():
     # Create a temporary directory for the test database
     temp_dir = tempfile.mkdtemp()
     try:
-        vector_store = LanceDBVectorStore(collection_name="test_collection")
+        vector_store = LanceDBVectorStore(
+            vector_store_schema_config=VectorStoreSchemaConfig(
+                index_name="test_collection"
+            )
+        )
         vector_store.connect(db_uri=temp_dir)
 
         docs = [
@@ -42,7 +47,7 @@ def test_vector_store_operations():
         ]
         vector_store.load_documents(docs[:2])
 
-        assert vector_store.collection_name in vector_store.db_connection.table_names()
+        assert vector_store.index_name in vector_store.db_connection.table_names()
 
         doc = vector_store.search_by_id("1")
         assert doc.id == "1"
@@ -91,7 +96,11 @@ def test_empty_collection():
     # Create a temporary directory for the test database
     temp_dir = tempfile.mkdtemp()
     try:
-        vector_store = LanceDBVectorStore(collection_name="empty_collection")
+        vector_store = LanceDBVectorStore(
+            vector_store_schema_config=VectorStoreSchemaConfig(
+                index_name="empty_collection"
+            )
+        )
         vector_store.connect(db_uri=temp_dir)
 
         # Load the vector store with a document, then delete it
@@ -102,12 +111,12 @@ def test_empty_collection():
             attributes={"title": "Tmp"},
         )
         vector_store.load_documents([sample_doc])
-        vector_store.db_connection.open_table(vector_store.collection_name).delete(
+        vector_store.db_connection.open_table(vector_store.index_name).delete(
             "id = 'tmp'"
         )
 
         # Should still have the collection
-        assert vector_store.collection_name in vector_store.db_connection.table_names()
+        assert vector_store.index_name in vector_store.db_connection.table_names()
 
         # Add a document after creating an empty collection
         doc = VectorStoreDocument(
@@ -131,7 +140,12 @@ def test_filter_search():
     # Create a temporary directory for the test database
     temp_dir = tempfile.mkdtemp()
     try:
-        vector_store = LanceDBVectorStore(collection_name="filter_collection")
+        vector_store = LanceDBVectorStore(
+            vector_store_schema_config=VectorStoreSchemaConfig(
+                index_name="filter_collection"
+            )
+        )
+
         vector_store.connect(db_uri=temp_dir)
 
         # Create test documents with different categories

tests/unit/query/context_builder/test_entity_extraction.py

@@ -3,6 +3,7 @@
 
 from typing import Any
 
+from graphrag.config.models.vector_store_schema_config import VectorStoreSchemaConfig
 from graphrag.data_model.entity import Entity
 from graphrag.data_model.types import TextEmbedder
 from graphrag.language_model.manager import ModelManager
@@ -19,7 +20,9 @@
 
 class MockBaseVectorStore(BaseVectorStore):
     def __init__(self, documents: list[VectorStoreDocument]) -> None:
-        super().__init__("mock")
+        super().__init__(
+            vector_store_schema_config=VectorStoreSchemaConfig(index_name="mock")
+        )
         self.documents = documents
 
     def connect(self, **kwargs: Any) -> None: