Add integration tests for WccCypherEndpoints

Author: DarthMax

graphdatascience/tests/integrationV2/procedure_surface/arrow/test_wcc_arrow_endpoints.py

@@ -25,8 +25,7 @@ def sample_graph(arrow_client: AuthenticatedArrowClient):
     (a)-[:REL]->(c)
     """
 
-    res = arrow_client.do_action( "v2/graph.fromGDL", json.dumps({"graphName": "g", "gdlGraph": gdl}).encode("utf-8"))
-    print(deserialize_single(res))
+    arrow_client.do_action( "v2/graph.fromGDL", json.dumps({"graphName": "g", "gdlGraph": gdl}).encode("utf-8"))
     yield MockGraph("g")
     arrow_client.do_action( "v2/graph.drop", json.dumps({"graphName": "g"}).encode("utf-8"))
 

graphdatascience/tests/integrationV2/procedure_surface/cypher/conftest.py

@@ -0,0 +1,34 @@
+import os
+from typing import Generator
+
+import pytest
+from testcontainers.core.container import DockerContainer
+from testcontainers.core.waiting_utils import wait_for_logs
+from testcontainers.neo4j import Neo4jContainer
+
+from graphdatascience.query_runner.neo4j_query_runner import Neo4jQueryRunner
+
+
+@pytest.fixture(scope="session")
+def neo4j_database_container() -> Generator[Neo4jContainer, None, None]:
+    neo4j_image = os.getenv("NEO4J_DATABASE_IMAGE", "neo4j:enterprise")
+
+    neo4j_container = (
+        Neo4jContainer(
+            image=neo4j_image,
+        )
+        .with_env("NEO4J_ACCEPT_LICENSE_AGREEMENT", "yes")
+        .with_env("NEO4J_PLUGINS", '["graph-data-science"]')
+    )
+
+    with neo4j_container as neo4j_db:
+        wait_for_logs(neo4j_db, "Started.")
+        yield neo4j_db
+
+
+@pytest.fixture
+def query_runner(neo4j_database_container: DockerContainer):
+    yield Neo4jQueryRunner.create_for_db(
+        f"bolt://localhost:{neo4j_database_container.get_exposed_port(7687)}",
+        ("neo4j", "password"),
+    )

graphdatascience/tests/integrationV2/procedure_surface/cypher/test_wcc_cypher_endpoints.py

@@ -0,0 +1,73 @@
+import pytest
+
+from graphdatascience import Graph, QueryRunner
+from graphdatascience.procedure_surface.arrow.arrow_wcc_endpoints import WccArrowEndpoints
+from graphdatascience.procedure_surface.cypher.wcc_cypher_endpoints import WccCypherEndpoints
+from graphdatascience.tests.integrationV2.procedure_surface.cypher.conftest import query_runner
+
+
+@pytest.fixture
+def sample_graph(query_runner: QueryRunner):
+    create_statement = """
+    CREATE
+    (a: Node),
+    (b: Node),
+    (c: Node),
+    (a)-[:REL]->(c)
+    """
+
+    query_runner.run_cypher(create_statement)
+
+    query_runner.run_cypher("""
+        MATCH (n)
+        OPTIONAL MATCH (n)-[r]->(m)
+        WITH gds.graph.project('g', n, m, {}) AS G
+        RETURN G
+    """)
+
+    yield Graph("g", query_runner)
+
+    query_runner.run_cypher("CALL gds.graph.drop('g')")
+    query_runner.run_cypher("MATCH (n) DETACH DELETE n")
+
+@pytest.fixture
+def wcc_endpoints(query_runner: QueryRunner):
+    yield WccCypherEndpoints(query_runner)
+
+
+def test_wcc_stats(wcc_endpoints: WccArrowEndpoints, sample_graph: Graph):
+    """Test WCC stats operation."""
+    result = wcc_endpoints.stats(
+        G=sample_graph
+    )
+
+    assert result.component_count == 2
+    assert result.compute_millis > 0
+    assert result.pre_processing_millis > 0
+    assert result.post_processing_millis > 0
+    assert "p10" in result.component_distribution
+
+def test_wcc_stream(wcc_endpoints: WccArrowEndpoints, sample_graph: Graph):
+    """Test WCC stream operation."""
+    result_df = wcc_endpoints.stream(
+        G=sample_graph,
+    )
+
+    assert "nodeId" in result_df.columns
+    assert "componentId" in result_df.columns
+    assert len(result_df.columns) == 2
+
+def test_wcc_mutate(wcc_endpoints: WccArrowEndpoints, sample_graph: Graph):
+    """Test WCC mutate operation."""
+    result = wcc_endpoints.mutate(
+        G=sample_graph,
+        mutate_property="componentId",
+    )
+
+    assert result.component_count == 2
+    assert "p10" in result.component_distribution
+    assert result.pre_processing_millis >= 0
+    assert result.compute_millis >= 0
+    assert result.post_processing_millis >= 0
+    assert result.mutate_millis >= 0
+    assert result.node_properties_written == 3