Add missing write tests for arrow endpoints

Author: DarthMax

graphdatascience/procedure_surface/arrow/celf_arrow_endpoints.py

@@ -147,7 +147,12 @@ def write(
         )
 
         result = self._node_property_endpoints.run_job_and_write(
-            "v2/centrality.celf", G, config, write_concurrency=write_concurrency, concurrency=concurrency, property_overwrites=write_property
+            "v2/centrality.celf",
+            G,
+            config,
+            write_concurrency=write_concurrency,
+            concurrency=concurrency,
+            property_overwrites=write_property,
         )
 
         return CelfWriteResult(**result)

graphdatascience/procedure_surface/arrow/graphsage_predict_arrow_endpoints.py

@@ -11,14 +11,15 @@
 )
 
 from ...arrow_client.authenticated_flight_client import AuthenticatedArrowClient
+from ...arrow_client.v2.remote_write_back_client import RemoteWriteBackClient
 from .model_api_arrow import ModelApiArrow
 from .node_property_endpoints import NodePropertyEndpoints
 
 
 class GraphSagePredictArrowEndpoints(GraphSagePredictEndpoints):
-    def __init__(self, arrow_client: AuthenticatedArrowClient):
+    def __init__(self, arrow_client: AuthenticatedArrowClient, write_back_client: Optional[RemoteWriteBackClient]):
         self._arrow_client = arrow_client
-        self._node_property_endpoints = NodePropertyEndpoints(arrow_client)
+        self._node_property_endpoints = NodePropertyEndpoints(arrow_client, write_back_client)
         self._model_api = ModelApiArrow(arrow_client)
 
     def stream(
@@ -79,12 +80,7 @@ def write(
         )
 
         raw_result = self._node_property_endpoints.run_job_and_write(
-            "v2/embeddings.graphSage",
-            G,
-            config,
-            write_concurrency,
-            concurrency,
-            write_property
+            "v2/embeddings.graphSage", G, config, write_concurrency, concurrency, write_property
         )
 
         return GraphSageWriteResult(**raw_result)

graphdatascience/procedure_surface/arrow/graphsage_train_arrow_endpoints.py

@@ -5,6 +5,8 @@
 from graphdatascience.procedure_surface.arrow.graphsage_predict_arrow_endpoints import GraphSagePredictArrowEndpoints
 
 from ...arrow_client.authenticated_flight_client import AuthenticatedArrowClient
+from ...arrow_client.v2.remote_write_back_client import RemoteWriteBackClient
+from ...graph.graph_object import Graph
 from ..api.graphsage_train_endpoints import (
     GraphSageTrainEndpoints,
     GraphSageTrainResult,
@@ -14,9 +16,10 @@
 
 
 class GraphSageTrainArrowEndpoints(GraphSageTrainEndpoints):
-    def __init__(self, arrow_client: AuthenticatedArrowClient):
+    def __init__(self, arrow_client: AuthenticatedArrowClient, write_back_client: Optional[RemoteWriteBackClient]):
         self._arrow_client = arrow_client
-        self._node_property_endpoints = NodePropertyEndpoints(arrow_client)
+        self._write_back_client = write_back_client
+        self._node_property_endpoints = NodePropertyEndpoints(arrow_client, write_back_client=write_back_client)
         self._model_api = ModelApiArrow(arrow_client)
 
     def train(
@@ -83,7 +86,9 @@ def train(
         result = self._node_property_endpoints.run_job_and_get_summary("v2/embeddings.graphSage.train", G, config)
 
         model = GraphSageModelV2(
-            model_name, self._model_api, predict_endpoints=GraphSagePredictArrowEndpoints(self._arrow_client)
+            model_name,
+            self._model_api,
+            predict_endpoints=GraphSagePredictArrowEndpoints(self._arrow_client, self._write_back_client),
         )
         train_result = GraphSageTrainResult(**result)
 

graphdatascience/session/session_v2_endpoints.py

@@ -68,11 +68,11 @@ def fast_rp(self) -> FastRPArrowEndpoints:
 
     @property
     def graphsage_predict(self) -> GraphSagePredictArrowEndpoints:
-        return GraphSagePredictArrowEndpoints(self._arrow_client)
+        return GraphSagePredictArrowEndpoints(self._arrow_client, self._write_back_client)
 
     @property
     def graphsage_train(self) -> GraphSageTrainArrowEndpoints:
-        return GraphSageTrainArrowEndpoints(self._arrow_client)
+        return GraphSageTrainArrowEndpoints(self._arrow_client, self._write_back_client)
 
     @property
     def harmonic_centrality(self) -> ClosenessHarmonicArrowEndpoints:

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

@@ -43,12 +43,14 @@ def create_graph_from_db(
     graph_name: str,
     graph_data: str,
     query: str,
+    undirected_relationship_types: Optional[list[str]] = None,
 ) -> Generator[GraphV2, Any, None]:
     try:
         query_runner.run_cypher(graph_data)
         result = CatalogArrowEndpoints(arrow_client, query_runner).project(
             graph_name=graph_name,
             query=query,
+            undirected_relationship_types=undirected_relationship_types,
         )
 
         yield result.graph

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

@@ -7,9 +7,10 @@
 from graphdatascience.arrow_client.v2.remote_write_back_client import RemoteWriteBackClient
 from graphdatascience.procedure_surface.api.catalog.graph_api import GraphV2
 from graphdatascience.procedure_surface.arrow.articlerank_arrow_endpoints import ArticleRankArrowEndpoints
-from graphdatascience.tests.integrationV2.procedure_surface.arrow.graph_creation_helper import create_graph, \
-    create_graph_from_db
-
+from graphdatascience.tests.integrationV2.procedure_surface.arrow.graph_creation_helper import (
+    create_graph,
+    create_graph_from_db,
+)
 
 graph = """
         CREATE
@@ -20,26 +21,29 @@
             (b)-[:REL]->(c)
         """
 
+
 @pytest.fixture
 def sample_graph(arrow_client: AuthenticatedArrowClient) -> Generator[GraphV2, None, None]:
     with create_graph(arrow_client, "g", graph) as G:
         yield G
 
+
 @pytest.fixture
 def db_graph(arrow_client: AuthenticatedArrowClient, query_runner: QueryRunner) -> Generator[Graph, None, None]:
     with create_graph_from_db(
-                arrow_client,
-                query_runner,
-                "g",
-                graph,
-                """
+        arrow_client,
+        query_runner,
+        "g",
+        graph,
+        """
                     MATCH (n)-->(m)
                     WITH gds.graph.project.remote(n, m) as g
                     RETURN g
-                """
+                """,
     ) as g:
         yield g
 
+
 @pytest.fixture
 def articlerank_endpoints(arrow_client: AuthenticatedArrowClient) -> Generator[ArticleRankArrowEndpoints, None, None]:
     yield ArticleRankArrowEndpoints(arrow_client)
@@ -85,8 +89,9 @@ def test_articlerank_mutate(articlerank_endpoints: ArticleRankArrowEndpoints, sa
     assert result.mutate_millis >= 0
     assert result.node_properties_written == 3
 
+
 def test_articlerank_write(arrow_client: AuthenticatedArrowClient, query_runner: QueryRunner, db_graph: Graph) -> None:
-    endpoints = ArticleRankArrowEndpoints(arrow_client, RemoteWriteBackClient(arrow_client,query_runner))
+    endpoints = ArticleRankArrowEndpoints(arrow_client, RemoteWriteBackClient(arrow_client, query_runner))
     result = endpoints.write(G=db_graph, write_property="write")
 
     assert result.did_converge

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

@@ -2,32 +2,56 @@
 
 import pytest
 
+from graphdatascience import Graph, QueryRunner
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
+from graphdatascience.arrow_client.v2.remote_write_back_client import RemoteWriteBackClient
 from graphdatascience.procedure_surface.api.articulationpoints_endpoints import (
     ArticulationPointsMutateResult,
     ArticulationPointsStatsResult,
+    ArticulationPointsWriteResult,
 )
 from graphdatascience.procedure_surface.api.catalog.graph_api import GraphV2
 from graphdatascience.procedure_surface.arrow.articulationpoints_arrow_endpoints import (
     ArticulationPointsArrowEndpoints,
 )
-from graphdatascience.tests.integrationV2.procedure_surface.arrow.graph_creation_helper import create_graph
-
+from graphdatascience.tests.integrationV2.procedure_surface.arrow.graph_creation_helper import (
+    create_graph,
+    create_graph_from_db,
+)
 
-@pytest.fixture
-def sample_graph(arrow_client: AuthenticatedArrowClient) -> Generator[GraphV2, None, None]:
-    gdl = """
-    (a: Node)
-    (b: Node)
-    (c: Node)
-    (a)-[:REL]->(c)
+graph = """
+    CREATE
+    (a: Node),
+    (b: Node),
+    (c: Node),
+    (a)-[:REL]->(c),
     (b)-[:REL]->(c)
     """
 
-    with create_graph(arrow_client, "g", gdl, undirected=("REL", "UNDIRECTED_REL")) as G:
+
+@pytest.fixture
+def sample_graph(arrow_client: AuthenticatedArrowClient) -> Generator[GraphV2, None, None]:
+    with create_graph(arrow_client, "g", graph, undirected=("REL", "UNDIRECTED_REL")) as G:
         yield G
 
 
+@pytest.fixture
+def db_graph(arrow_client: AuthenticatedArrowClient, query_runner: QueryRunner) -> Generator[Graph, None, None]:
+    with create_graph_from_db(
+        arrow_client,
+        query_runner,
+        "g",
+        graph,
+        """
+                    MATCH (n)-->(m)
+                    WITH gds.graph.project.remote(n, m, {relationshipType: "REL"}) as g
+                    RETURN g
+                """,
+        ["REL"],
+    ) as g:
+        yield g
+
+
 @pytest.fixture
 def articulationpoints_endpoints(arrow_client: AuthenticatedArrowClient) -> ArticulationPointsArrowEndpoints:
     return ArticulationPointsArrowEndpoints(arrow_client)
@@ -36,7 +60,6 @@ def articulationpoints_endpoints(arrow_client: AuthenticatedArrowClient) -> Arti
 def test_articulationpoints_mutate(
     articulationpoints_endpoints: ArticulationPointsArrowEndpoints, sample_graph: GraphV2
 ) -> None:
-    """Test ArticulationPoints mutate operation."""
     result = articulationpoints_endpoints.mutate(
         G=sample_graph,
         mutate_property="articulationPoint",
@@ -53,7 +76,6 @@ def test_articulationpoints_mutate(
 def test_articulationpoints_stats(
     articulationpoints_endpoints: ArticulationPointsArrowEndpoints, sample_graph: GraphV2
 ) -> None:
-    """Test ArticulationPoints stats operation."""
     result = articulationpoints_endpoints.stats(sample_graph)
 
     assert isinstance(result, ArticulationPointsStatsResult)
@@ -64,17 +86,33 @@ def test_articulationpoints_stats(
 def test_articulationpoints_stream_not_implemented(
     articulationpoints_endpoints: ArticulationPointsArrowEndpoints, sample_graph: GraphV2
 ) -> None:
-    """Test that ArticulationPoints stream raises NotImplementedError."""
     with pytest.raises(
         NotImplementedError, match="Stream mode is not supported for ArticulationPoints arrow endpoints"
     ):
         articulationpoints_endpoints.stream(sample_graph)
 
 
+def test_articulationpoints_write(
+    arrow_client: AuthenticatedArrowClient, query_runner: QueryRunner, db_graph: Graph
+) -> None:
+    endpoints = ArticulationPointsArrowEndpoints(arrow_client, RemoteWriteBackClient(arrow_client, query_runner))
+    result = endpoints.write(G=db_graph, write_property="articulationPoint")
+
+    assert isinstance(result, ArticulationPointsWriteResult)
+    assert result.articulation_point_count >= 0
+    assert result.compute_millis >= 0
+    assert result.write_millis >= 0
+    assert result.node_properties_written == 3
+
+    assert (
+        query_runner.run_cypher("MATCH (n) WHERE n.articulationPoint IS NOT NULL RETURN COUNT(*) AS count").squeeze()
+        == 3
+    )
+
+
 def test_articulationpoints_estimate(
     articulationpoints_endpoints: ArticulationPointsArrowEndpoints, sample_graph: GraphV2
 ) -> None:
-    """Test ArticulationPoints memory estimation."""
     result = articulationpoints_endpoints.estimate(sample_graph)
 
     assert result.node_count == 3