Add MutationClient

Author: DarthMax

graphdatascience/arrow_client/v2/mutation_client.py

@@ -0,0 +1,16 @@
+import json
+
+from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
+from graphdatascience.arrow_client.v2.api_types import MutateResult
+from graphdatascience.arrow_client.v2.data_mapper_utils import deserialize_single
+
+
+class MutationClient:
+    MUTATE_ENDPOINT = "v2/results.mutate"
+
+    @staticmethod
+    def mutate_node_property(client: AuthenticatedArrowClient, job_id: str, mutate_property: str) -> MutateResult:
+        mutate_config = {"jobId": job_id, "mutateProperty": mutate_property}
+        encoded_config = json.dumps(mutate_config).encode("utf-8")
+        mutate_arrow_res = client.do_action_with_retry(MutationClient.MUTATE_ENDPOINT, encoded_config)
+        return MutateResult(**deserialize_single(mutate_arrow_res))

graphdatascience/tests/unit/arrow_client/V2/test_job_client.py

@@ -12,17 +12,14 @@ def setUp(self) -> None:
         self.mock_client = MagicMock()
 
     def test_run_job(self) -> None:
-        # Setup
         job_id = "test-job-123"
         endpoint = "v2/test.endpoint"
         config = {"param1": "value1", "param2": 42}
 
         self.mock_client.do_action_with_retry.return_value = iter([ArrowTestResult({"jobId": job_id})])
 
-        # Execute
         result = JobClient.run_job(self.mock_client, endpoint, config)
 
-        # Verify
         expected_config = json.dumps(config).encode("utf-8")
         self.mock_client.do_action_with_retry.assert_called_once_with(endpoint, expected_config)
         self.assertEqual(result, job_id)