From 7cecb9218d249c2bb0776bc9163e8852059a874d Mon Sep 17 00:00:00 2001
From: Matthew Roeschke <10647082+mroeschke@users.noreply.github.com>
Date: Mon, 29 Apr 2024 13:48:20 -0700
Subject: [PATCH 1/3] Remove return in test_lbfgs

---
 python/cuml/tests/dask/test_dask_logistic_regression.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/python/cuml/tests/dask/test_dask_logistic_regression.py b/python/cuml/tests/dask/test_dask_logistic_regression.py
index 4fe7504ce4..f6c68238ae 100644
--- a/python/cuml/tests/dask/test_dask_logistic_regression.py
+++ b/python/cuml/tests/dask/test_dask_logistic_regression.py
@@ -425,8 +425,6 @@ def array_to_numpy(ary):
         np.abs(accuracy_cuml - accuracy_sk) < 1e-3
     )
 
-    return lr
-
 
 @pytest.mark.parametrize("fit_intercept", [False, True])
 def test_noreg(fit_intercept, client):

From c738bdf2b088a523a79225c8597817f39386559d Mon Sep 17 00:00:00 2001
From: Matthew Roeschke <10647082+mroeschke@users.noreply.github.com>
Date: Mon, 29 Apr 2024 17:20:42 -0700
Subject: [PATCH 2/3] Add back return but make private

---
 .../dask/test_dask_logistic_regression.py     | 38 ++++++++++++-------
 1 file changed, 24 insertions(+), 14 deletions(-)

diff --git a/python/cuml/tests/dask/test_dask_logistic_regression.py b/python/cuml/tests/dask/test_dask_logistic_regression.py
index f6c68238ae..556c02c8f1 100644
--- a/python/cuml/tests/dask/test_dask_logistic_regression.py
+++ b/python/cuml/tests/dask/test_dask_logistic_regression.py
@@ -303,14 +303,7 @@ def assert_params(
     )
 
 
-@pytest.mark.mg
-@pytest.mark.parametrize("nrows", [1e5])
-@pytest.mark.parametrize("ncols", [20])
-@pytest.mark.parametrize("n_parts", [2, 23])
-@pytest.mark.parametrize("fit_intercept", [False, True])
-@pytest.mark.parametrize("datatype", [np.float32])
-@pytest.mark.parametrize("delayed", [True, False])
-def test_lbfgs(
+def _test_lbfgs(
     nrows,
     ncols,
     n_parts,
@@ -425,10 +418,27 @@ def array_to_numpy(ary):
         np.abs(accuracy_cuml - accuracy_sk) < 1e-3
     )
 
+    return lr
+
+
+@pytest.mark.mg
+@pytest.mark.parametrize("fit_intercept", [False, True])
+@pytest.mark.parametrize("delayed", [True, False])
+def test_lbfgs(n_parts, fit_intercept, delayed, client):
+    _test_lbfgs(
+        nrows=1e5,
+        ncols=20,
+        n_parts=n_parts,
+        fit_intercept=fit_intercept,
+        datatype=np.float32,
+        delayed=delayed,
+        client=client,
+    )
+
 
 @pytest.mark.parametrize("fit_intercept", [False, True])
 def test_noreg(fit_intercept, client):
-    lr = test_lbfgs(
+    lr = _test_lbfgs(
         nrows=1e5,
         ncols=20,
         n_parts=23,
@@ -492,7 +502,7 @@ def assert_small(X, y, n_classes):
 @pytest.mark.parametrize("n_classes", [8])
 def test_n_classes(n_parts, fit_intercept, n_classes, client):
     nrows = int(1e5) if n_classes < 5 else int(2e5)
-    lr = test_lbfgs(
+    lr = _test_lbfgs(
         nrows=nrows,
         ncols=20,
         n_parts=n_parts,
@@ -515,7 +525,7 @@ def test_n_classes(n_parts, fit_intercept, n_classes, client):
 @pytest.mark.parametrize("C", [1.0, 10.0])
 def test_l1(fit_intercept, datatype, delayed, n_classes, C, client):
     nrows = int(1e5) if n_classes < 5 else int(2e5)
-    lr = test_lbfgs(
+    lr = _test_lbfgs(
         nrows=nrows,
         ncols=20,
         n_parts=2,
@@ -543,7 +553,7 @@ def test_elasticnet(
     fit_intercept, datatype, delayed, n_classes, l1_ratio, client
 ):
     nrows = int(1e5) if n_classes < 5 else int(2e5)
-    lr = test_lbfgs(
+    lr = _test_lbfgs(
         nrows=nrows,
         ncols=20,
         n_parts=2,
@@ -583,7 +593,7 @@ def test_sparse_from_dense(
 
     nrows = int(1e5) if n_classes < 5 else int(2e5)
     run_test = partial(
-        test_lbfgs,
+        _test_lbfgs,
         nrows=nrows,
         ncols=20,
         n_parts=2,
@@ -697,7 +707,7 @@ def test_standardization_on_normal_dataset(
     nrows = int(1e5) if n_classes < 5 else int(2e5)
 
     # test correctness compared with scikit-learn
-    test_lbfgs(
+    _test_lbfgs(
         nrows=nrows,
         ncols=20,
         n_parts=2,

From 58dfaaebe6ff7206cb17e20ba82f4d71684d9535 Mon Sep 17 00:00:00 2001
From: Matthew Roeschke <10647082+mroeschke@users.noreply.github.com>
Date: Tue, 30 Apr 2024 10:19:04 -0700
Subject: [PATCH 3/3] Add back param

---
 python/cuml/tests/dask/test_dask_logistic_regression.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/python/cuml/tests/dask/test_dask_logistic_regression.py b/python/cuml/tests/dask/test_dask_logistic_regression.py
index 556c02c8f1..a512d78d4f 100644
--- a/python/cuml/tests/dask/test_dask_logistic_regression.py
+++ b/python/cuml/tests/dask/test_dask_logistic_regression.py
@@ -422,6 +422,7 @@ def array_to_numpy(ary):
 
 
 @pytest.mark.mg
+@pytest.mark.parametrize("n_parts", [2, 23])
 @pytest.mark.parametrize("fit_intercept", [False, True])
 @pytest.mark.parametrize("delayed", [True, False])
 def test_lbfgs(n_parts, fit_intercept, delayed, client):