diff --git a/python/cugraph/cugraph/tests/utils/test_replication_mg.py b/python/cugraph/cugraph/tests/utils/test_replication_mg.py index 2f9c0d0189b..13089bc238b 100644 --- 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)