From f1490b0b3a6014bf8635560af6f745e15f9c4a77 Mon Sep 17 00:00:00 2001
From: Ralph Liu <137829296+nv-rliu@users.noreply.github.com>
Date: Mon, 20 May 2024 10:47:48 -0400
Subject: [PATCH] Fix MG Betweenness Centrality Test Bug (#4429)

Fixes bug seen in MG tests due to incorrect call to `enable_batch`

Authors:
  - Ralph Liu (https://github.com/nv-rliu)

Approvers:
  - Alex Barghi (https://github.com/alexbarghi-nv)

URL: https://github.com/rapidsai/cugraph/pull/4429
---
 .../centrality/test_betweenness_centrality.py     | 15 +++++----------
 1 file changed, 5 insertions(+), 10 deletions(-)

diff --git a/python/cugraph/cugraph/tests/centrality/test_betweenness_centrality.py b/python/cugraph/cugraph/tests/centrality/test_betweenness_centrality.py
index ff8859a01b1..6d1f53f7fc3 100644
--- a/python/cugraph/cugraph/tests/centrality/test_betweenness_centrality.py
+++ b/python/cugraph/cugraph/tests/centrality/test_betweenness_centrality.py
@@ -111,18 +111,13 @@ def calc_betweenness_centrality(
     else:
         edge_attr = None
 
-    G = None
+    G = graph_file.get_graph(
+        download=True,
+        create_using=cugraph.Graph(directed=directed),
+        ignore_weights=not edgevals,
+    )
     if multi_gpu_batch:
-        G = graph_file.get_dask_graph(
-            create_using=cugraph.Graph(directed=directed), ignore_weights=not edgevals
-        )
         G.enable_batch()
-    else:
-        G = graph_file.get_graph(
-            download=True,
-            create_using=cugraph.Graph(directed=directed),
-            ignore_weights=not edgevals,
-        )
 
     M = G.to_pandas_edgelist().rename(
         columns={"src": "0", "dst": "1", "wgt": edge_attr}