Fix CI

Author: stellasia

tests/e2e/experimental/test_simplekgpipeline_e2e.py

@@ -128,9 +128,7 @@ async def test_pipeline_builder_happy_path_legacy_schema(
 
     # check the content of the graph:
     # check lexical graph content
-    records, _, _ = driver.execute_query(
-        "MATCH (start:chunkNodeLabel) RETURN start"
-    )
+    records, _, _ = driver.execute_query("MATCH (start:chunkNodeLabel) RETURN start")
     assert len(records) == 1
 
     # check entity -> chunk relationships
@@ -162,8 +160,8 @@ async def test_pipeline_builder_happy_path(
 
     # Instantiate schema
     entities = ["Person"]
-    relations = []
-    potential_schema = []
+    relations: list[str] = []
+    potential_schema: list[tuple[str, str, str]] = []
     schema = {
         "node_types": entities,
         "relationship_types": relations,
@@ -181,7 +179,7 @@ async def test_pipeline_builder_happy_path(
         llm=llm,
         driver=driver,
         embedder=embedder,
-        schema=schema,
+        schema=schema,  # type: ignore[arg-type]
         from_pdf=from_pdf,
         on_error=on_error,
         lexical_graph_config=lexical_graph_config,
@@ -192,9 +190,7 @@ async def test_pipeline_builder_happy_path(
 
     # check the content of the graph:
     # check lexical graph content
-    records, _, _ = driver.execute_query(
-        "MATCH (start:chunkNodeLabel) RETURN start"
-    )
+    records, _, _ = driver.execute_query("MATCH (start:chunkNodeLabel) RETURN start")
     assert len(records) == 1
 
     # check entity -> chunk relationships

tests/unit/experimental/components/test_graph_pruning.py

@@ -167,13 +167,15 @@ def test_graph_pruning_validate_node(
         assert result is None
 
 
-def test_graph_pruning_enforce_nodes_lexical_graph(lexical_graph_config: LexicalGraphConfig) -> None:
+def test_graph_pruning_enforce_nodes_lexical_graph(
+    lexical_graph_config: LexicalGraphConfig,
+) -> None:
     pruner = GraphPruning()
     result = pruner._enforce_nodes(
         extracted_nodes=[
             Neo4jNode(id="1", label="Paragraph"),
         ],
-        schema=GraphSchema(additional_node_types=False),
+        schema=GraphSchema(node_types=tuple(), additional_node_types=False),
         lexical_graph_config=lexical_graph_config,
         pruning_stats=PruningStats(),
     )

tests/unit/experimental/pipeline/test_kg_builder.py

@@ -18,10 +18,6 @@
 import neo4j
 import pytest
 from neo4j_graphrag.embeddings import Embedder
-from neo4j_graphrag.experimental.components.schema import (
-    NodeType,
-    RelationshipType,
-)
 from neo4j_graphrag.experimental.components.types import LexicalGraphConfig
 from neo4j_graphrag.experimental.pipeline.exceptions import PipelineDefinitionError
 from neo4j_graphrag.experimental.pipeline.kg_builder import SimpleKGPipeline

tests/unit/test_schema.py

@@ -38,7 +38,7 @@
 
 
 def _query_return_value(*args: Any, **kwargs: Any) -> list[Any]:
-    query = kwargs.get("query", args[1] if len(args) > 1 else None)
+    query = kwargs.get("query", args[1] if len(args) > 1 else "")
     if NODE_PROPERTIES_QUERY in query:
         return [
             {