Add integration tests for WccArrowEndpoints

Author: DarthMax

graphdatascience/arrow_client/v2/job_client.py

@@ -45,19 +45,24 @@ def get_summary(client: AuthenticatedArrowClient, job_id: str) -> dict[str, Any]
         return deserialize_single(res)
 
     @staticmethod
-    def stream_results(client: AuthenticatedArrowClient, job_id: str) -> DataFrame:
-        encoded_config = JobIdConfig(jobId=job_id).dump_json().encode("utf-8")
+    def stream_results(client: AuthenticatedArrowClient,graph_name: str, job_id: str) -> DataFrame:
+        payload = {
+            "graphName": graph_name,
+            "jobId": job_id,
+        }
 
-        res = client.do_action_with_retry("v2/results.stream", encoded_config)
+        res = client.do_action_with_retry("v2/results.stream", json.dumps(payload).encode("utf-8"))
         export_job_id = JobIdConfig(**deserialize_single(res)).job_id
 
         payload = {
+            "version": "v2",
             "name": export_job_id,
-            "version": 1,
+            "body": {}
         }
 
         ticket = Ticket(json.dumps(payload).encode("utf-8"))
-        with client.get_stream(ticket) as get:
-            arrow_table = get.read_all()
 
+        get = client.get_stream(ticket)
+        arrow_table = get.read_all()
         return arrow_table.to_pandas(types_mapper=ArrowDtype)  # type: ignore
+

graphdatascience/procedure_surface/arrow/arrow_wcc_endpoints.py

@@ -1,13 +1,13 @@
-from typing import List, Optional, Any
+from typing import Any, List, Optional
 
 from pandas import DataFrame
 
-from ..api.estimation_result import EstimationResult
 from ...arrow_client.authenticated_flight_client import AuthenticatedArrowClient
 from ...arrow_client.v2.job_client import JobClient
 from ...arrow_client.v2.mutation_client import MutationClient
 from ...arrow_client.v2.write_back_client import WriteBackClient
 from ...graph.graph_object import Graph
+from ..api.estimation_result import EstimationResult
 from ..api.wcc_endpoints import WccEndpoints, WccMutateResult, WccStatsResult, WccWriteResult
 from ..utils.config_converter import ConfigConverter
 

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

@@ -0,0 +1,73 @@
+import json
+
+import pytest
+
+from graphdatascience import Graph
+from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
+from graphdatascience.arrow_client.v2.data_mapper_utils import deserialize_single
+from graphdatascience.procedure_surface.arrow.arrow_wcc_endpoints import WccArrowEndpoints
+
+
+class MockGraph(Graph):
+    def __init__(self, name: str):
+        self._name = name
+
+    def name(self) -> str:
+        return self._name
+
+
+@pytest.fixture
+def sample_graph(arrow_client: AuthenticatedArrowClient):
+    gdl = """
+    (a: Node)
+    (b: Node)
+    (c: Node)
+    (a)-[:REL]->(c)
+    """
+
+    res = arrow_client.do_action( "v2/graph.fromGDL", json.dumps({"graphName": "g", "gdlGraph": gdl}).encode("utf-8"))
+    print(deserialize_single(res))
+    yield MockGraph("g")
+    arrow_client.do_action( "v2/graph.drop", json.dumps({"graphName": "g"}).encode("utf-8"))
+
+@pytest.fixture
+def wcc_endpoints(arrow_client: AuthenticatedArrowClient):
+    yield WccArrowEndpoints(arrow_client)
+
+
+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