Merge pull request #957 from DarthMax/implement_missing_write_tests

GDSA-146 Add missing write tests

Author: DarthMax

graphdatascience/procedure_surface/arrow/articlerank_arrow_endpoints.py

@@ -172,7 +172,7 @@ def write(
         )
 
         result = self._node_property_endpoints.run_job_and_write(
-            "v2/centrality.articleRank", G, config, write_concurrency, concurrency
+            "v2/centrality.articleRank", G, config, write_concurrency, concurrency, write_property
         )
 
         return ArticleRankWriteResult(**result)

graphdatascience/procedure_surface/arrow/articulationpoints_arrow_endpoints.py

@@ -123,7 +123,7 @@ def write(
         )
 
         result = self._node_property_endpoints.run_job_and_write(
-            "v2/centrality.articulationPoints", G, config, write_concurrency, concurrency
+            "v2/centrality.articulationPoints", G, config, write_concurrency, concurrency, write_property
         )
 
         return ArticulationPointsWriteResult(**result)

graphdatascience/procedure_surface/arrow/betweenness_arrow_endpoints.py

@@ -154,7 +154,7 @@ def write(
         )
 
         result = self._node_property_endpoints.run_job_and_write(
-            "v2/centrality.betweenness", G, config, write_concurrency, concurrency
+            "v2/centrality.betweenness", G, config, write_concurrency, concurrency, write_property
         )
 
         return BetweennessWriteResult(**result)

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
+            "v2/centrality.celf",
+            G,
+            config,
+            write_concurrency=write_concurrency,
+            concurrency=concurrency,
+            property_overwrites=write_property,
         )
 
         return CelfWriteResult(**result)

graphdatascience/procedure_surface/arrow/closeness_arrow_endpoints.py

@@ -134,7 +134,7 @@ def write(
         )
 
         result = self._node_property_endpoints.run_job_and_write(
-            "v2/centrality.closeness", G, config, write_concurrency, concurrency
+            "v2/centrality.closeness", G, config, write_concurrency, concurrency, write_property
         )
 
         return ClosenessWriteResult(**result)

graphdatascience/procedure_surface/arrow/closeness_harmonic_arrow_endpoints.py

@@ -124,7 +124,7 @@ def write(
         )
 
         result = self._node_property_endpoints.run_job_and_write(
-            "v2/centrality.harmonic", G, config, write_concurrency, concurrency
+            "v2/centrality.harmonic", G, config, write_concurrency, concurrency, write_property
         )
 
         return ClosenessHarmonicWriteResult(**result)

graphdatascience/procedure_surface/arrow/degree_arrow_endpoints.py

@@ -131,7 +131,7 @@ def write(
         )
 
         result = self._node_property_endpoints.run_job_and_write(
-            "v2/centrality.degree", G, config, write_concurrency, concurrency
+            "v2/centrality.degree", G, config, write_concurrency, concurrency, write_property
         )
 
         return DegreeWriteResult(**result)

graphdatascience/procedure_surface/arrow/eigenvector_arrow_endpoints.py

@@ -168,7 +168,7 @@ def write(
         )
 
         result = self._node_property_endpoints.run_job_and_write(
-            "v2/centrality.eigenvector", G, config, write_concurrency, concurrency
+            "v2/centrality.eigenvector", G, config, write_concurrency, concurrency, write_property
         )
 
         return EigenvectorWriteResult(**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,11 +80,7 @@ def write(
         )
 
         raw_result = self._node_property_endpoints.run_job_and_write(
-            "v2/embeddings.graphSage",
-            G,
-            config,
-            write_concurrency,
-            concurrency,
+            "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,7 @@
 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 ..api.graphsage_train_endpoints import (
     GraphSageTrainEndpoints,
     GraphSageTrainResult,
@@ -14,9 +15,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 +85,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)