Add a write test for ArticleRankArrowEndpoints

Author: DarthMax

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

@@ -2,25 +2,43 @@
 
 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.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
+from graphdatascience.tests.integrationV2.procedure_surface.arrow.graph_creation_helper import create_graph, \
+    create_graph_from_db
 
 
+graph = """
+        CREATE
+            (a: Node),
+            (b: Node),
+            (c: Node),
+            (a)-[:REL]->(c),
+            (b)-[:REL]->(c)
+        """
+
 @pytest.fixture
 def sample_graph(arrow_client: AuthenticatedArrowClient) -> Generator[GraphV2, None, None]:
-    gdl = """
-    (a: Node)
-    (b: Node)
-    (c: Node)
-    (a)-[:REL]->(c)
-    (b)-[:REL]->(c)
-    """
-
-    with create_graph(arrow_client, "g", gdl) as G:
+    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,
+                """
+                    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]:
@@ -67,6 +85,20 @@ 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))
+    result = endpoints.write(G=db_graph, write_property="write")
+
+    assert result.did_converge
+    assert "p50" in result.centrality_distribution
+    assert result.pre_processing_millis >= 0
+    assert result.compute_millis >= 0
+    assert result.post_processing_millis >= 0
+    assert result.write_millis >= 0
+    assert result.node_properties_written == 3
+
+    assert query_runner.run_cypher("MATCH (n) WHERE n.write IS NOT NULL RETURN COUNT(*) AS count").squeeze() == 3
+
 
 def test_articlerank_estimate(articlerank_endpoints: ArticleRankArrowEndpoints, sample_graph: GraphV2) -> None:
     result = articlerank_endpoints.estimate(sample_graph)