Update replicated edgelist test after using new API

rapidsai · Jun 4, 2024 · 4a44f2f · 4a44f2f
1 parent 826b86b
commit 4a44f2f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/python/cugraph/cugraph/tests/utils/test_replication_mg.py b/python/cugraph/cugraph/tests/utils/test_replication_mg.py
@@ -200,8 +200,8 @@ def test_enable_batch_edgelist_replication(graph_file, directed, dask_client):
     G = utils.generate_cugraph_graph_from_file(graph_file, directed)
     G.enable_batch()
     df = G.edgelist.edgelist_df
-    for worker in G.batch_edgelists:
-        replicated_df = G.batch_edgelists[worker].result()
+    for i in range(G.batch_edgelists.npartitions):
+        replicated_df = G.batch_edgelists.get_partition(i).compute()
         assert_frame_equal(df, replicated_df)