cognee 0.4.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

cognee/tests/test_relational_db_migration.py

@@ -1,6 +1,5 @@
 import pathlib
 import os
-from typing import List
 from cognee.infrastructure.databases.graph import get_graph_engine
 from cognee.infrastructure.databases.relational import (
     get_migration_relational_engine,
@@ -10,7 +9,7 @@ from cognee.infrastructure.databases.vector.pgvector import (
     create_db_and_tables as create_pgvector_db_and_tables,
 )
 from cognee.tasks.ingestion import migrate_relational_database
-from cognee.modules.search.types import SearchResult, SearchType
+from cognee.modules.search.types import SearchType
 import cognee
 
 
@@ -27,6 +26,9 @@ def normalize_node_name(node_name: str) -> str:
 
 
 async def setup_test_db():
+    # Disable backend access control to migrate relational data
+    os.environ["ENABLE_BACKEND_ACCESS_CONTROL"] = "false"
+
     await cognee.prune.prune_data()
     await cognee.prune.prune_system(metadata=True)
 
@@ -271,6 +273,55 @@ async def test_schema_only_migration():
     print(f"Edge counts: {edge_counts}")
 
 
+async def test_search_result_quality():
+    from cognee.infrastructure.databases.relational import (
+        get_migration_relational_engine,
+    )
+
+    # Get relational database with original data
+    migration_engine = get_migration_relational_engine()
+    from sqlalchemy import text
+
+    async with migration_engine.engine.connect() as conn:
+        result = await conn.execute(
+            text("""
+                SELECT
+                    c.CustomerId,
+                    c.FirstName,
+                    c.LastName,
+                    GROUP_CONCAT(i.InvoiceId, ',') AS invoice_ids
+                FROM Customer AS c
+                LEFT JOIN Invoice AS i ON c.CustomerId = i.CustomerId
+                GROUP BY c.CustomerId, c.FirstName, c.LastName
+            """)
+        )
+
+        for row in result:
+            # Get expected invoice IDs from relational DB for each Customer
+            customer_id = row.CustomerId
+            invoice_ids = row.invoice_ids.split(",") if row.invoice_ids else []
+            print(f"Relational DB Customer {customer_id}: {invoice_ids}")
+
+            # Use Cognee search to get invoice IDs for the same Customer but by providing Customer name
+            search_results = await cognee.search(
+                query_type=SearchType.GRAPH_COMPLETION,
+                query_text=f"List me all the invoices of Customer:{row.FirstName} {row.LastName}.",
+                top_k=50,
+                system_prompt="Just return me the invoiceID as a number without any text. This is an example output: ['1', '2', '3']. Where 1, 2, 3 are invoiceIDs of an invoice",
+            )
+            print(f"Cognee search result: {search_results}")
+
+            import ast
+
+            lst = ast.literal_eval(search_results[0])  # converts string -> Python list
+            # Transfrom both lists to int for comparison, sorting and type consistency
+            lst = sorted([int(x) for x in lst])
+            invoice_ids = sorted([int(x) for x in invoice_ids])
+            assert lst == invoice_ids, (
+                f"Search results {lst} do not match expected invoice IDs {invoice_ids} for Customer:{customer_id}"
+            )
+
+
 async def test_migration_sqlite():
     database_to_migrate_path = os.path.join(pathlib.Path(__file__).parent, "test_data/")
 
@@ -283,6 +334,7 @@ async def test_migration_sqlite():
     )
 
     await relational_db_migration()
+    await test_search_result_quality()
     await test_schema_only_migration()
 
 

cognee/tests/test_search_db.py

@@ -2,6 +2,7 @@ import pathlib
 import os
 import cognee
 from cognee.infrastructure.databases.graph import get_graph_engine
+from cognee.infrastructure.databases.vector import get_vector_engine
 from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
 from cognee.modules.graph.utils import resolve_edges_to_text
 from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever
@@ -12,8 +13,10 @@ from cognee.modules.retrieval.graph_completion_cot_retriever import GraphComplet
 from cognee.modules.retrieval.graph_summary_completion_retriever import (
     GraphSummaryCompletionRetriever,
 )
+from cognee.modules.retrieval.triplet_retriever import TripletRetriever
 from cognee.shared.logging_utils import get_logger
 from cognee.modules.search.types import SearchType
+from cognee.modules.users.methods import get_default_user
 from collections import Counter
 
 logger = get_logger()
@@ -37,6 +40,23 @@ async def main():
 
     await cognee.cognify([dataset_name])
 
+    user = await get_default_user()
+    from cognee.memify_pipelines.create_triplet_embeddings import create_triplet_embeddings
+
+    await create_triplet_embeddings(user=user, dataset=dataset_name, triplets_batch_size=5)
+
+    graph_engine = await get_graph_engine()
+    nodes, edges = await graph_engine.get_graph_data()
+
+    vector_engine = get_vector_engine()
+    collection = await vector_engine.search(
+        query_text="Test", limit=None, collection_name="Triplet_text"
+    )
+
+    assert len(edges) == len(collection), (
+        f"Expected {len(edges)} edges but got {len(collection)} in Triplet_text collection"
+    )
+
     context_gk = await GraphCompletionRetriever().get_context(
         query="Next to which country is Germany located?"
     )
@@ -49,6 +69,9 @@ async def main():
     context_gk_sum = await GraphSummaryCompletionRetriever().get_context(
         query="Next to which country is Germany located?"
     )
+    context_triplet = await TripletRetriever().get_context(
+        query="Next to which country is Germany located?"
+    )
 
     for name, context in [
         ("GraphCompletionRetriever", context_gk),
@@ -65,6 +88,13 @@ async def main():
             f"{name}: Context did not contain 'germany' or 'netherlands'; got: {context!r}"
         )
 
+    assert isinstance(context_triplet, str), "TripletRetriever: Context should be a string"
+    assert len(context_triplet) > 0, "TripletRetriever: Context should not be empty"
+    lower_triplet = context_triplet.lower()
+    assert "germany" in lower_triplet or "netherlands" in lower_triplet, (
+        f"TripletRetriever: Context did not contain 'germany' or 'netherlands'; got: {context_triplet!r}"
+    )
+
     triplets_gk = await GraphCompletionRetriever().get_triplets(
         query="Next to which country is Germany located?"
     )
@@ -129,6 +159,11 @@ async def main():
         query_text="Next to which country is Germany located?",
         save_interaction=True,
     )
+    completion_triplet = await cognee.search(
+        query_type=SearchType.TRIPLET_COMPLETION,
+        query_text="Next to which country is Germany located?",
+        save_interaction=True,
+    )
 
     await cognee.search(
         query_type=SearchType.FEEDBACK,
@@ -141,12 +176,19 @@ async def main():
         ("GRAPH_COMPLETION_COT", completion_cot),
         ("GRAPH_COMPLETION_CONTEXT_EXTENSION", completion_ext),
         ("GRAPH_SUMMARY_COMPLETION", completion_sum),
+        ("TRIPLET_COMPLETION", completion_triplet),
     ]:
         assert isinstance(search_results, list), f"{name}: should return a list"
         assert len(search_results) == 1, (
             f"{name}: expected single-element list, got {len(search_results)}"
         )
-        text = search_results[0]
+
+        from cognee.context_global_variables import backend_access_control_enabled
+
+        if backend_access_control_enabled():
+            text = search_results[0]["search_result"][0]
+        else:
+            text = search_results[0]
         assert isinstance(text, str), f"{name}: element should be a string"
         assert text.strip(), f"{name}: string should not be empty"
         assert "netherlands" in text.lower(), (
@@ -162,7 +204,7 @@ async def main():
 
     # Assert there are exactly 4 CogneeUserInteraction nodes.
     assert type_counts.get("CogneeUserInteraction", 0) == 4, (
-        f"Expected exactly four DCogneeUserInteraction nodes, but found {type_counts.get('CogneeUserInteraction', 0)}"
+        f"Expected exactly four CogneeUserInteraction nodes, but found {type_counts.get('CogneeUserInteraction', 0)}"
     )
 
     # Assert there is exactly two CogneeUserFeedback nodes.

cognee/tests/unit/api/test_conditional_authentication_endpoints.py

@@ -1,3 +1,4 @@
+import os
 import pytest
 from unittest.mock import patch, AsyncMock, MagicMock
 from uuid import uuid4
@@ -5,8 +6,6 @@ from fastapi.testclient import TestClient
 from types import SimpleNamespace
 import importlib
 
-from cognee.api.client import app
-
 
 # Fixtures for reuse across test classes
 @pytest.fixture
@@ -32,6 +31,10 @@ def mock_authenticated_user():
     )
 
 
+# To turn off authentication we need to set the environment variable before importing the module
+# Also both require_authentication and backend access control must be false
+os.environ["REQUIRE_AUTHENTICATION"] = "false"
+os.environ["ENABLE_BACKEND_ACCESS_CONTROL"] = "false"
 gau_mod = importlib.import_module("cognee.modules.users.methods.get_authenticated_user")
 
 
@@ -40,6 +43,8 @@ class TestConditionalAuthenticationEndpoints:
 
     @pytest.fixture
     def client(self):
+        from cognee.api.client import app
+
         """Create a test client."""
         return TestClient(app)
 
@@ -133,6 +138,8 @@ class TestConditionalAuthenticationBehavior:
 
     @pytest.fixture
     def client(self):
+        from cognee.api.client import app
+
         return TestClient(app)
 
     @pytest.mark.parametrize(
@@ -209,6 +216,8 @@ class TestConditionalAuthenticationErrorHandling:
 
     @pytest.fixture
     def client(self):
+        from cognee.api.client import app
+
         return TestClient(app)
 
     @patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
@@ -232,7 +241,7 @@ class TestConditionalAuthenticationErrorHandling:
         # The exact error message may vary depending on the actual database connection
         # The important thing is that we get a 500 error when user creation fails
 
-    def test_current_environment_configuration(self):
+    def test_current_environment_configuration(self, client):
         """Test that current environment configuration is working properly."""
         # This tests the actual module state without trying to change it
         from cognee.modules.users.methods.get_authenticated_user import (

cognee/tests/unit/api/test_ontology_endpoint.py

@@ -0,0 +1,252 @@
+import pytest
+import uuid
+from fastapi.testclient import TestClient
+from unittest.mock import Mock
+from types import SimpleNamespace
+from cognee.api.client import app
+from cognee.modules.users.methods import get_authenticated_user
+
+
+@pytest.fixture(scope="session")
+def test_client():
+    # Keep a single TestClient (and event loop) for the whole module.
+    # Re-creating TestClient repeatedly can break async DB connections (asyncpg loop mismatch).
+    with TestClient(app) as c:
+        yield c
+
+
+@pytest.fixture
+def client(test_client, mock_default_user):
+    async def override_get_authenticated_user():
+        return mock_default_user
+
+    app.dependency_overrides[get_authenticated_user] = override_get_authenticated_user
+    yield test_client
+    app.dependency_overrides.pop(get_authenticated_user, None)
+
+
+@pytest.fixture
+def mock_user():
+    user = Mock()
+    user.id = "test-user-123"
+    return user
+
+
+@pytest.fixture
+def mock_default_user():
+    """Mock default user for testing."""
+    return SimpleNamespace(
+        id=str(uuid.uuid4()),
+        email="default@example.com",
+        is_active=True,
+        tenant_id=str(uuid.uuid4()),
+    )
+
+
+def test_upload_ontology_success(client):
+    """Test successful ontology upload"""
+    ontology_content = (
+        b"<rdf:RDF xmlns:rdf='http://www.w3.org/1999/02/22-rdf-syntax-ns#'></rdf:RDF>"
+    )
+    unique_key = f"test_ontology_{uuid.uuid4().hex[:8]}"
+
+    response = client.post(
+        "/api/v1/ontologies",
+        files=[("ontology_file", ("test.owl", ontology_content, "application/xml"))],
+        data={"ontology_key": unique_key, "description": "Test"},
+    )
+
+    assert response.status_code == 200
+    data = response.json()
+    assert data["uploaded_ontologies"][0]["ontology_key"] == unique_key
+    assert "uploaded_at" in data["uploaded_ontologies"][0]
+
+
+def test_upload_ontology_invalid_file(client):
+    """Test 400 response for non-.owl files"""
+    unique_key = f"test_ontology_{uuid.uuid4().hex[:8]}"
+    response = client.post(
+        "/api/v1/ontologies",
+        files={"ontology_file": ("test.txt", b"not xml")},
+        data={"ontology_key": unique_key},
+    )
+    assert response.status_code == 400
+
+
+def test_upload_ontology_missing_data(client):
+    """Test 400 response for missing file or key"""
+    # Missing file
+    response = client.post("/api/v1/ontologies", data={"ontology_key": "test"})
+    assert response.status_code == 400
+
+    # Missing key
+    response = client.post(
+        "/api/v1/ontologies", files=[("ontology_file", ("test.owl", b"xml", "application/xml"))]
+    )
+    assert response.status_code == 400
+
+
+def test_upload_ontology_without_auth_header(client):
+    """Test behavior when no explicit authentication header is provided."""
+    unique_key = f"test_ontology_{uuid.uuid4().hex[:8]}"
+    response = client.post(
+        "/api/v1/ontologies",
+        files=[("ontology_file", ("test.owl", b"<rdf></rdf>", "application/xml"))],
+        data={"ontology_key": unique_key},
+    )
+
+    assert response.status_code == 200
+    data = response.json()
+    assert data["uploaded_ontologies"][0]["ontology_key"] == unique_key
+    assert "uploaded_at" in data["uploaded_ontologies"][0]
+
+
+def test_upload_multiple_ontologies_in_single_request_is_rejected(client):
+    """Uploading multiple ontology files in a single request should fail."""
+    import io
+
+    file1_content = b"<rdf:RDF xmlns:rdf='http://www.w3.org/1999/02/22-rdf-syntax-ns#'></rdf:RDF>"
+    file2_content = b"<rdf:RDF xmlns:rdf='http://www.w3.org/1999/02/22-rdf-syntax-ns#'></rdf:RDF>"
+
+    files = [
+        ("ontology_file", ("vehicles.owl", io.BytesIO(file1_content), "application/xml")),
+        ("ontology_file", ("manufacturers.owl", io.BytesIO(file2_content), "application/xml")),
+    ]
+    data = {"ontology_key": "vehicles", "description": "Base vehicles"}
+
+    response = client.post("/api/v1/ontologies", files=files, data=data)
+
+    assert response.status_code == 400
+    assert "Only one ontology_file is allowed" in response.json()["error"]
+
+
+def test_upload_endpoint_rejects_array_style_fields(client):
+    """Array-style form values should be rejected (no backwards compatibility)."""
+    import io
+    import json
+
+    file_content = b"<rdf:RDF xmlns:rdf='http://www.w3.org/1999/02/22-rdf-syntax-ns#'></rdf:RDF>"
+
+    files = [("ontology_file", ("single.owl", io.BytesIO(file_content), "application/xml"))]
+    data = {
+        "ontology_key": json.dumps(["single_key"]),
+        "description": json.dumps(["Single ontology"]),
+    }
+
+    response = client.post("/api/v1/ontologies", files=files, data=data)
+
+    assert response.status_code == 400
+    assert "ontology_key must be a string" in response.json()["error"]
+
+
+def test_cognify_with_multiple_ontologies(client):
+    """Test cognify endpoint accepts multiple ontology keys"""
+    payload = {
+        "datasets": ["test_dataset"],
+        "ontology_key": ["ontology1", "ontology2"],  # Array instead of string
+        "run_in_background": False,
+    }
+
+    response = client.post("/api/v1/cognify", json=payload)
+
+    # Should not fail due to ontology_key type
+    assert response.status_code in [200, 400, 409]  # May fail for other reasons, not type
+
+
+def test_complete_multifile_workflow(client):
+    """Test workflow: upload ontologies one-by-one → cognify with multiple keys"""
+    import io
+
+    # Step 1: Upload two ontologies (one-by-one)
+    file1_content = b"""<?xml version="1.0"?>
+    <rdf:RDF xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#"
+             xmlns:owl="http://www.w3.org/2002/07/owl#">
+        <owl:Class rdf:ID="Vehicle"/>
+    </rdf:RDF>"""
+
+    file2_content = b"""<?xml version="1.0"?>
+    <rdf:RDF xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#"
+             xmlns:owl="http://www.w3.org/2002/07/owl#">
+        <owl:Class rdf:ID="Manufacturer"/>
+    </rdf:RDF>"""
+
+    upload_response_1 = client.post(
+        "/api/v1/ontologies",
+        files=[("ontology_file", ("vehicles.owl", io.BytesIO(file1_content), "application/xml"))],
+        data={"ontology_key": "vehicles", "description": "Vehicle ontology"},
+    )
+    assert upload_response_1.status_code == 200
+
+    upload_response_2 = client.post(
+        "/api/v1/ontologies",
+        files=[
+            ("ontology_file", ("manufacturers.owl", io.BytesIO(file2_content), "application/xml"))
+        ],
+        data={"ontology_key": "manufacturers", "description": "Manufacturer ontology"},
+    )
+    assert upload_response_2.status_code == 200
+
+    # Step 2: Verify ontologies are listed
+    list_response = client.get("/api/v1/ontologies")
+    assert list_response.status_code == 200
+    ontologies = list_response.json()
+    assert "vehicles" in ontologies
+    assert "manufacturers" in ontologies
+
+    # Step 3: Test cognify with multiple ontologies
+    cognify_payload = {
+        "datasets": ["test_dataset"],
+        "ontology_key": ["vehicles", "manufacturers"],
+        "run_in_background": False,
+    }
+
+    cognify_response = client.post("/api/v1/cognify", json=cognify_payload)
+    # Should not fail due to ontology handling (may fail for dataset reasons)
+    assert cognify_response.status_code != 400  # Not a validation error
+
+
+def test_upload_error_handling(client):
+    """Test error handling for invalid uploads (single-file endpoint)."""
+    import io
+    import json
+
+    # Array-style key should be rejected
+    file_content = b"<rdf:RDF></rdf:RDF>"
+    files = [("ontology_file", ("test.owl", io.BytesIO(file_content), "application/xml"))]
+    data = {
+        "ontology_key": json.dumps(["key1", "key2"]),
+        "description": "desc1",
+    }
+
+    response = client.post("/api/v1/ontologies", files=files, data=data)
+    assert response.status_code == 400
+    assert "ontology_key must be a string" in response.json()["error"]
+
+    # Duplicate key should be rejected
+    response_1 = client.post(
+        "/api/v1/ontologies",
+        files=[("ontology_file", ("test1.owl", io.BytesIO(file_content), "application/xml"))],
+        data={"ontology_key": "duplicate", "description": "desc1"},
+    )
+    assert response_1.status_code == 200
+
+    response_2 = client.post(
+        "/api/v1/ontologies",
+        files=[("ontology_file", ("test2.owl", io.BytesIO(file_content), "application/xml"))],
+        data={"ontology_key": "duplicate", "description": "desc2"},
+    )
+    assert response_2.status_code == 400
+    assert "already exists" in response_2.json()["error"]
+
+
+def test_cognify_missing_ontology_key(client):
+    """Test cognify with non-existent ontology key"""
+    payload = {
+        "datasets": ["test_dataset"],
+        "ontology_key": ["nonexistent_key"],
+        "run_in_background": False,
+    }
+
+    response = client.post("/api/v1/cognify", json=payload)
+    assert response.status_code == 409
+    assert "Ontology key 'nonexistent_key' not found" in response.json()["error"]

cognee/tests/unit/infrastructure/databases/cache/test_cache_config.py

@@ -8,6 +8,7 @@ def test_cache_config_defaults():
     """Test that CacheConfig has the correct default values."""
     config = CacheConfig()
 
+    assert config.cache_backend == "fs"
     assert config.caching is False
     assert config.shared_kuzu_lock is False
     assert config.cache_host == "localhost"
@@ -19,6 +20,7 @@ def test_cache_config_defaults():
 def test_cache_config_custom_values():
     """Test that CacheConfig accepts custom values."""
     config = CacheConfig(
+        cache_backend="redis",
         caching=True,
         shared_kuzu_lock=True,
         cache_host="redis.example.com",
@@ -27,6 +29,7 @@ def test_cache_config_custom_values():
         agentic_lock_timeout=180,
     )
 
+    assert config.cache_backend == "redis"
     assert config.caching is True
     assert config.shared_kuzu_lock is True
     assert config.cache_host == "redis.example.com"
@@ -38,6 +41,7 @@ def test_cache_config_custom_values():
 def test_cache_config_to_dict():
     """Test the to_dict method returns all configuration values."""
     config = CacheConfig(
+        cache_backend="fs",
         caching=True,
         shared_kuzu_lock=True,
         cache_host="test-host",
@@ -49,6 +53,7 @@ def test_cache_config_to_dict():
     config_dict = config.to_dict()
 
     assert config_dict == {
+        "cache_backend": "fs",
         "caching": True,
         "shared_kuzu_lock": True,
         "cache_host": "test-host",

cognee/tests/unit/infrastructure/databases/test_index_data_points.py

@@ -0,0 +1,27 @@
+import pytest
+from unittest.mock import AsyncMock, patch, MagicMock
+from cognee.tasks.storage.index_data_points import index_data_points
+from cognee.infrastructure.engine import DataPoint
+
+
+class TestDataPoint(DataPoint):
+    name: str
+    metadata: dict = {"index_fields": ["name"]}
+
+
+@pytest.mark.asyncio
+async def test_index_data_points_calls_vector_engine():
+    """Test that index_data_points creates vector index and indexes data."""
+    data_points = [TestDataPoint(name="test1")]
+
+    mock_vector_engine = AsyncMock()
+    mock_vector_engine.embedding_engine.get_batch_size = MagicMock(return_value=100)
+
+    with patch.dict(
+        index_data_points.__globals__,
+        {"get_vector_engine": lambda: mock_vector_engine},
+    ):
+        await index_data_points(data_points)
+
+    assert mock_vector_engine.create_vector_index.await_count >= 1
+    assert mock_vector_engine.index_data_points.await_count >= 1

cognee/tests/unit/infrastructure/databases/test_index_graph_edges.py

@@ -5,8 +5,7 @@ from cognee.tasks.storage.index_graph_edges import index_graph_edges
 
 @pytest.mark.asyncio
 async def test_index_graph_edges_success():
-    """Test that index_graph_edges uses the index datapoints and creates vector index."""
-    # Create the mocks for the graph and vector engines.
+    """Test that index_graph_edges retrieves edges and delegates to index_data_points."""
     mock_graph_engine = AsyncMock()
     mock_graph_engine.get_graph_data.return_value = (
         None,
@@ -15,26 +14,23 @@ async def test_index_graph_edges_success():
             [{"relationship_name": "rel2"}],
         ],
     )
-    mock_vector_engine = AsyncMock()
-    mock_vector_engine.embedding_engine.get_batch_size = MagicMock(return_value=100)
+    mock_index_data_points = AsyncMock()
 
-    # Patch the globals of the function so that when it does:
-    #   vector_engine = get_vector_engine()
-    #   graph_engine = await get_graph_engine()
-    # it uses the mocked versions.
     with patch.dict(
         index_graph_edges.__globals__,
         {
             "get_graph_engine": AsyncMock(return_value=mock_graph_engine),
-            "get_vector_engine": lambda: mock_vector_engine,
+            "index_data_points": mock_index_data_points,
         },
     ):
         await index_graph_edges()
 
-    # Assertions on the mock calls.
     mock_graph_engine.get_graph_data.assert_awaited_once()
-    assert mock_vector_engine.create_vector_index.await_count == 1
-    assert mock_vector_engine.index_data_points.await_count == 1
+    mock_index_data_points.assert_awaited_once()
+
+    call_args = mock_index_data_points.call_args[0][0]
+    assert len(call_args) == 2
+    assert all(hasattr(item, "relationship_name") for item in call_args)
 
 
 @pytest.mark.asyncio
@@ -42,20 +38,22 @@ async def test_index_graph_edges_no_relationships():
     """Test that index_graph_edges handles empty relationships correctly."""
     mock_graph_engine = AsyncMock()
     mock_graph_engine.get_graph_data.return_value = (None, [])
-    mock_vector_engine = AsyncMock()
+    mock_index_data_points = AsyncMock()
 
     with patch.dict(
         index_graph_edges.__globals__,
         {
             "get_graph_engine": AsyncMock(return_value=mock_graph_engine),
-            "get_vector_engine": lambda: mock_vector_engine,
+            "index_data_points": mock_index_data_points,
         },
     ):
         await index_graph_edges()
 
     mock_graph_engine.get_graph_data.assert_awaited_once()
-    mock_vector_engine.create_vector_index.assert_not_awaited()
-    mock_vector_engine.index_data_points.assert_not_awaited()
+    mock_index_data_points.assert_awaited_once()
+
+    call_args = mock_index_data_points.call_args[0][0]
+    assert len(call_args) == 0
 
 
 @pytest.mark.asyncio