Use pytest conventions consistently (#6337)

* Do not derive from unittest.TestCase (not needed for pytest)

* assertRaises -> pytest.raises

* Simplify test_empty_dmatrix with test parametrization

* setUpClass -> setup_class, tearDownClass -> teardown_class

* Don't import unittest; import pytest

* Use plain assert

* Use parametrized tests in more places

* Fix test_gpu_with_sklearn.py

* Put back run_empty_dmatrix_reg / run_empty_dmatrix_cls

* Fix test_eta_decay_gpu_hist

* Add parametrized tests for monotone constraints

* Fix test names

* Remove test parametrization

* Revise test_slice to be not flaky

tests/python-gpu/load_pickle.py

@@ -1,6 +1,5 @@
 '''Loading a pickled model generated by test_pickling.py, only used by
 `test_gpu_with_dask.py`'''
-import unittest
 import os
 import numpy as np
 import xgboost as xgb
@@ -14,7 +13,7 @@ sys.path.append("tests/python")
 import testing as tm
 
 
-class TestLoadPickle(unittest.TestCase):
+class TestLoadPickle:
     def test_load_pkl(self):
         '''Test whether prediction is correct.'''
         assert os.environ['CUDA_VISIBLE_DEVICES'] == '-1'

tests/python-gpu/test_device_quantile_dmatrix.py

@@ -1,7 +1,6 @@
 # -*- coding: utf-8 -*-
 import numpy as np
 import xgboost as xgb
-import unittest
 import pytest
 import sys
 
@@ -9,7 +8,7 @@ sys.path.append("tests/python")
 import testing as tm
 
 
-class TestDeviceQuantileDMatrix(unittest.TestCase):
+class TestDeviceQuantileDMatrix:
     def test_dmatrix_numpy_init(self):
         data = np.random.randn(5, 5)
         with pytest.raises(TypeError,

tests/python-gpu/test_gpu_basic_models.py

@@ -1,15 +1,15 @@
 import sys
 import os
-import unittest
 import numpy as np
 import xgboost as xgb
+import pytest
 sys.path.append("tests/python")
 # Don't import the test class, otherwise they will run twice.
 import test_callback as test_cb  # noqa
 rng = np.random.RandomState(1994)
 
 
-class TestGPUBasicModels(unittest.TestCase):
+class TestGPUBasicModels:
     cputest = test_cb.TestCallbacks()
 
     def run_cls(self, X, y, deterministic):

tests/python-gpu/test_gpu_interaction_constraints.py

@@ -1,5 +1,4 @@
 import numpy as np
-import unittest
 import sys
 sys.path.append("tests/python")
 # Don't import the test class, otherwise they will run twice.
@@ -7,7 +6,7 @@ import test_interaction_constraints as test_ic  # noqa
 rng = np.random.RandomState(1994)
 
 
-class TestGPUInteractionConstraints(unittest.TestCase):
+class TestGPUInteractionConstraints:
     cputest = test_ic.TestInteractionConstraints()
 
     def test_interaction_constraints(self):

tests/python-gpu/test_gpu_pickling.py

@@ -1,6 +1,5 @@
 '''Test model IO with pickle.'''
 import pickle
-import unittest
 import numpy as np
 import subprocess
 import os
@@ -35,7 +34,7 @@ def load_pickle(path):
     return bst
 
 
-class TestPickling(unittest.TestCase):
+class TestPickling:
     args_template = [
         "pytest",
         "--verbose",

tests/python-gpu/test_gpu_prediction.py

@@ -1,6 +1,4 @@
 import sys
-import json
-import unittest
 import pytest
 
 import numpy as np
@@ -26,7 +24,7 @@ predict_parameter_strategy = strategies.fixed_dictionaries({
 })
 
 
-class TestGPUPredict(unittest.TestCase):
+class TestGPUPredict:
     def test_predict(self):
         iterations = 10
         np.random.seed(1)

tests/python-gpu/test_gpu_ranking.py

@@ -1,16 +1,15 @@
 import numpy as np
 import xgboost
 import os
-import unittest
 import itertools
 import shutil
 import urllib.request
 import zipfile
 
 
-class TestRanking(unittest.TestCase):
+class TestRanking:
     @classmethod
-    def setUpClass(cls):
+    def setup_class(cls):
         """
         Download and setup the test fixtures
         """
@@ -75,7 +74,7 @@ class TestRanking(unittest.TestCase):
                           'predictor': 'cpu_predictor'}
 
     @classmethod
-    def tearDownClass(cls):
+    def teardown_class(cls):
         """
         Cleanup test artifacts from download and unpacking
         :return:

tests/python-gpu/test_gpu_training_continuation.py

@@ -1,4 +1,3 @@
-import unittest
 import numpy as np
 import xgboost as xgb
 import json
@@ -6,7 +5,7 @@ import json
 rng = np.random.RandomState(1994)
 
 
-class TestGPUTrainingContinuation(unittest.TestCase):
+class TestGPUTrainingContinuation:
     def run_training_continuation(self, use_json):
         kRows = 64
         kCols = 32

tests/python-gpu/test_gpu_with_sklearn.py

@@ -2,7 +2,6 @@ import xgboost as xgb
 import pytest
 import sys
 import numpy as np
-import unittest
 
 sys.path.append("tests/python")
 import testing as tm               # noqa
@@ -33,8 +32,5 @@ def test_gpu_binary_classification():
             assert err < 0.1
 
 
-class TestGPUBoostFromPrediction(unittest.TestCase):
+def test_boost_from_prediction_gpu_hist():
-    cpu_test = twskl.TestBoostFromPrediction()
+    cpu_test = twskl.run_boost_from_prediction('gpu_hist')
-
-    def test_boost_from_prediction_gpu_hist(self):
-        self.cpu_test.run_boost_from_prediction('gpu_hist')

tests/python-gpu/test_monotonic_constraints.py

@@ -1,7 +1,6 @@
 import sys
 import numpy as np
 
-import unittest
 import pytest
 
 import xgboost as xgb
@@ -38,26 +37,27 @@ def assert_constraint(constraint, tree_method):
         assert non_increasing(pred)
 
 
-class TestMonotonicConstraints(unittest.TestCase):
+@pytest.mark.skipif(**tm.no_sklearn())
-    @pytest.mark.skipif(**tm.no_sklearn())
+def test_gpu_hist_basic():
-    def test_gpu_hist_basic(self):
+    assert_constraint(1, 'gpu_hist')
-        assert_constraint(1, 'gpu_hist')
+    assert_constraint(-1, 'gpu_hist')
-        assert_constraint(-1, 'gpu_hist')
 
-    def test_gpu_hist_depthwise(self):
-        params = {
-            'tree_method': 'gpu_hist',
-            'grow_policy': 'depthwise',
-            'monotone_constraints': '(1, -1)'
-        }
-        model = xgb.train(params, tmc.training_dset)
-        tmc.is_correctly_constrained(model)
 
-    def test_gpu_hist_lossguide(self):
+def test_gpu_hist_depthwise():
-        params = {
+    params = {
-            'tree_method': 'gpu_hist',
+        'tree_method': 'gpu_hist',
-            'grow_policy': 'lossguide',
+        'grow_policy': 'depthwise',
-            'monotone_constraints': '(1, -1)'
+        'monotone_constraints': '(1, -1)'
-        }
+    }
-        model = xgb.train(params, tmc.training_dset)
+    model = xgb.train(params, tmc.training_dset)
-        tmc.is_correctly_constrained(model)
+    tmc.is_correctly_constrained(model)
+
+
+def test_gpu_hist_lossguide():
+    params = {
+        'tree_method': 'gpu_hist',
+        'grow_policy': 'lossguide',
+        'monotone_constraints': '(1, -1)'
+    }
+    model = xgb.train(params, tmc.training_dset)
+    tmc.is_correctly_constrained(model)

tests/python/test_basic.py

@@ -2,7 +2,7 @@
 import numpy as np
 import os
 import xgboost as xgb
-import unittest
+import pytest
 import json
 from pathlib import Path
 import tempfile
@@ -12,7 +12,7 @@ dpath = 'demo/data/'
 rng = np.random.RandomState(1994)
 
 
-class TestBasic(unittest.TestCase):
+class TestBasic:
     def test_compat(self):
         from xgboost.compat import lazy_isinstance
         a = np.array([1, 2, 3])
@@ -119,30 +119,28 @@ class TestBasic(unittest.TestCase):
 
         # number of feature importances should == number of features
         dump1 = bst.get_dump()
-        self.assertEqual(len(dump1), 1, "Expected only 1 tree to be dumped.")
+        assert len(dump1) == 1, 'Expected only 1 tree to be dumped.'
-        self.assertEqual(len(dump1[0].splitlines()), 3,
+        len(dump1[0].splitlines()) == 3, 'Expected 1 root and 2 leaves - 3 lines in dump.'
-                         "Expected 1 root and 2 leaves - 3 lines in dump.")
 
         dump2 = bst.get_dump(with_stats=True)
-        self.assertEqual(dump2[0].count('\n'), 3,
+        assert dump2[0].count('\n') == 3, 'Expected 1 root and 2 leaves - 3 lines in dump.'
-                         "Expected 1 root and 2 leaves - 3 lines in dump.")
+        assert (dump2[0].find('\n') > dump1[0].find('\n'),
-        self.assertGreater(dump2[0].find('\n'), dump1[0].find('\n'),
+                'Expected more info when with_stats=True is given.')
-                           "Expected more info when with_stats=True is given.")
 
         dump3 = bst.get_dump(dump_format="json")
         dump3j = json.loads(dump3[0])
-        self.assertEqual(dump3j["nodeid"], 0, "Expected the root node on top.")
+        assert dump3j['nodeid'] == 0, 'Expected the root node on top.'
 
         dump4 = bst.get_dump(dump_format="json", with_stats=True)
         dump4j = json.loads(dump4[0])
-        self.assertIn("gain", dump4j, "Expected 'gain' to be dumped in JSON.")
+        assert 'gain' in dump4j, "Expected 'gain' to be dumped in JSON."
 
     def test_load_file_invalid(self):
-        self.assertRaises(xgb.core.XGBoostError, xgb.Booster,
+        with pytest.raises(xgb.core.XGBoostError):
-                          model_file='incorrect_path')
+            xgb.Booster(model_file='incorrect_path')
 
-        self.assertRaises(xgb.core.XGBoostError, xgb.Booster,
+        with pytest.raises(xgb.core.XGBoostError):
-                          model_file=u'不正なパス')
+            xgb.Booster(model_file=u'不正なパス')
 
     def test_dmatrix_numpy_init_omp(self):
 
@@ -226,7 +224,7 @@ class TestBasic(unittest.TestCase):
         assert output == solution
 
 
-class TestBasicPathLike(unittest.TestCase):
+class TestBasicPathLike:
     """Unit tests using pathlib.Path for file interaction."""
 
     def test_DMatrix_init_from_path(self):
@@ -253,8 +251,8 @@ class TestBasicPathLike(unittest.TestCase):
 
     def test_Booster_init_invalid_path(self):
         """An invalid model_file path should raise XGBoostError."""
-        self.assertRaises(xgb.core.XGBoostError, xgb.Booster,
+        with pytest.raises(xgb.core.XGBoostError):
-                          model_file=Path("invalidpath"))
+            xgb.Booster(model_file=Path("invalidpath"))
 
 
     def test_Booster_save_and_load(self):

tests/python/test_basic_models.py

@@ -1,6 +1,5 @@
 import numpy as np
 import xgboost as xgb
-import unittest
 import os
 import json
 import testing as tm

tests/python/test_callback.py

@@ -1,5 +1,4 @@
 import xgboost as xgb
-import unittest
 import pytest
 import os
 import testing as tm
@@ -9,9 +8,9 @@ import tempfile
 pytestmark = pytest.mark.skipif(**tm.no_sklearn())
 
 
-class TestCallbacks(unittest.TestCase):
+class TestCallbacks:
     @classmethod
-    def setUpClass(cls):
+    def setup_class(cls):
         from sklearn.datasets import load_breast_cancer
         X, y = load_breast_cancer(return_X_y=True)
         cls.X = X
@@ -148,9 +147,9 @@ class TestCallbacks(unittest.TestCase):
         early_stop = xgb.callback.EarlyStopping(rounds=early_stopping_rounds,
                                                 save_best=True)
         cls = xgb.XGBClassifier(booster='gblinear', n_estimators=10)
-        self.assertRaises(ValueError, lambda: cls.fit(X, y, eval_set=[(X, y)],
+        with pytest.raises(ValueError):
-                                                      eval_metric=tm.eval_error_metric,
+            cls.fit(X, y, eval_set=[(X, y)], eval_metric=tm.eval_error_metric,
-                                                      callbacks=[early_stop]))
+                    callbacks=[early_stop])
 
         # No error
         early_stop = xgb.callback.EarlyStopping(rounds=early_stopping_rounds,
@@ -172,17 +171,20 @@ class TestCallbacks(unittest.TestCase):
         watchlist = [(dtest, 'eval'), (dtrain, 'train')]
         num_round = 4
 
+        warning_check = pytest.warns(UserWarning) if deprecated_callback else tm.noop_context()
+
         # learning_rates as a list
         # init eta with 0 to check whether learning_rates work
         param = {'max_depth': 2, 'eta': 0, 'verbosity': 0,
                  'objective': 'binary:logistic', 'eval_metric': 'error',
                  'tree_method': tree_method}
         evals_result = {}
-        bst = xgb.train(param, dtrain, num_round, watchlist,
+        with warning_check:
-                        callbacks=[scheduler([
+            bst = xgb.train(param, dtrain, num_round, watchlist,
-                            0.8, 0.7, 0.6, 0.5
+                            callbacks=[scheduler([
-                        ])],
+                                0.8, 0.7, 0.6, 0.5
-                        evals_result=evals_result)
+                            ])],
+                            evals_result=evals_result)
         eval_errors_0 = list(map(float, evals_result['eval']['error']))
         assert isinstance(bst, xgb.core.Booster)
         # validation error should decrease, if eta > 0
@@ -193,10 +195,11 @@ class TestCallbacks(unittest.TestCase):
                  'objective': 'binary:logistic', 'eval_metric': 'error',
                  'tree_method': tree_method}
         evals_result = {}
-        bst = xgb.train(param, dtrain, num_round, watchlist,
+        with warning_check:
-                        callbacks=[scheduler(
+            bst = xgb.train(param, dtrain, num_round, watchlist,
-                            [0.8, 0.7, 0.6, 0.5])],
+                            callbacks=[scheduler(
-                        evals_result=evals_result)
+                                [0.8, 0.7, 0.6, 0.5])],
+                            evals_result=evals_result)
         eval_errors_1 = list(map(float, evals_result['eval']['error']))
         assert isinstance(bst, xgb.core.Booster)
         # validation error should decrease, if learning_rate > 0
@@ -208,11 +211,12 @@ class TestCallbacks(unittest.TestCase):
             'eval_metric': 'error', 'tree_method': tree_method
         }
         evals_result = {}
-        bst = xgb.train(param, dtrain, num_round, watchlist,
+        with warning_check:
-                        callbacks=[scheduler(
+            bst = xgb.train(param, dtrain, num_round, watchlist,
-                            [0, 0, 0, 0]
+                            callbacks=[scheduler(
-                        )],
+                                [0, 0, 0, 0]
-                        evals_result=evals_result)
+                            )],
+                            evals_result=evals_result)
         eval_errors_2 = list(map(float, evals_result['eval']['error']))
         assert isinstance(bst, xgb.core.Booster)
         # validation error should not decrease, if eta/learning_rate = 0
@@ -223,11 +227,12 @@ class TestCallbacks(unittest.TestCase):
             return num_boost_round / (ithround + 1)
 
         evals_result = {}
-        bst = xgb.train(param, dtrain, num_round, watchlist,
+        with warning_check:
-                        callbacks=[
+            bst = xgb.train(param, dtrain, num_round, watchlist,
-                            scheduler(eta_decay)
+                            callbacks=[
-                        ],
+                                scheduler(eta_decay)
-                        evals_result=evals_result)
+                            ],
+                            evals_result=evals_result)
         eval_errors_3 = list(map(float, evals_result['eval']['error']))
 
         assert isinstance(bst, xgb.core.Booster)
@@ -238,18 +243,15 @@ class TestCallbacks(unittest.TestCase):
             assert eval_errors_3[i] != eval_errors_2[i]
 
     def test_eta_decay_hist(self):
-        with pytest.warns(UserWarning):
+        self.run_eta_decay('hist', True)
-            self.run_eta_decay('hist', True)
         self.run_eta_decay('hist', False)
 
     def test_eta_decay_approx(self):
-        with pytest.warns(UserWarning):
+        self.run_eta_decay('approx', True)
-            self.run_eta_decay('approx', True)
         self.run_eta_decay('approx', False)
 
     def test_eta_decay_exact(self):
-        with pytest.warns(UserWarning):
+        self.run_eta_decay('exact', True)
-            self.run_eta_decay('exact', True)
         self.run_eta_decay('exact', False)
 
     def test_check_point(self):

tests/python/test_cli.py

@@ -1,6 +1,5 @@
 import os
 import tempfile
-import unittest
 import platform
 import xgboost
 import subprocess
@@ -9,7 +8,7 @@ import json
 import testing as tm
 
 
-class TestCLI(unittest.TestCase):
+class TestCLI:
     template = '''
 booster = gbtree
 objective = reg:squarederror

tests/python/test_dmatrix.py

@@ -1,7 +1,6 @@
 # -*- coding: utf-8 -*-
 import numpy as np
 import xgboost as xgb
-import unittest
 import scipy.sparse
 import pytest
 from scipy.sparse import rand, csr_matrix
@@ -12,7 +11,7 @@ dpath = 'demo/data/'
 rng = np.random.RandomState(1994)
 
 
-class TestDMatrix(unittest.TestCase):
+class TestDMatrix:
     def test_warn_missing(self):
         from xgboost import data
         with pytest.warns(UserWarning):
@@ -48,15 +47,19 @@ class TestDMatrix(unittest.TestCase):
         assert dm.num_col() == 2
 
         # 0d array
-        self.assertRaises(ValueError, xgb.DMatrix, np.array(1))
+        with pytest.raises(ValueError):
+            xgb.DMatrix(np.array(1))
         # 1d array
-        self.assertRaises(ValueError, xgb.DMatrix, np.array([1, 2, 3]))
+        with pytest.raises(ValueError):
+            xgb.DMatrix(np.array([1, 2, 3]))
         # 3d array
         data = np.random.randn(5, 5, 5)
-        self.assertRaises(ValueError, xgb.DMatrix, data)
+        with pytest.raises(ValueError):
+            xgb.DMatrix(data)
         # object dtype
         data = np.array([['a', 'b'], ['c', 'd']])
-        self.assertRaises(ValueError, xgb.DMatrix, data)
+        with pytest.raises(ValueError):
+            xgb.DMatrix(data)
 
     def test_csr(self):
         indptr = np.array([0, 2, 3, 6])
@@ -106,56 +109,50 @@ class TestDMatrix(unittest.TestCase):
 
     def test_slice(self):
         X = rng.randn(100, 100)
-        y = rng.randint(low=0, high=3, size=100)
+        y = rng.randint(low=0, high=3, size=100).astype(np.float32)
         d = xgb.DMatrix(X, y)
-        np.testing.assert_equal(d.get_label(), y.astype(np.float32))
+        np.testing.assert_equal(d.get_label(), y)
 
         fw = rng.uniform(size=100).astype(np.float32)
         d.set_info(feature_weights=fw)
 
-        eval_res_0 = {}
+        # base margin is per-class in multi-class classifier
-        booster = xgb.train(
+        base_margin = rng.randn(100, 3).astype(np.float32)
-            {'num_class': 3, 'objective': 'multi:softprob',
+        d.set_base_margin(base_margin.flatten())
-             'eval_metric': 'merror'},
-            d,
-            num_boost_round=2, evals=[(d, 'd')], evals_result=eval_res_0)
-
-        predt = booster.predict(d)
-        predt = predt.reshape(100 * 3, 1)
-
-        d.set_base_margin(predt)
 
         ridxs = [1, 2, 3, 4, 5, 6]
         sliced = d.slice(ridxs)
 
-        sliced_margin = sliced.get_float_info('base_margin')
+        # Slicing works with label and other meta info fields
-        assert sliced_margin.shape[0] == len(ridxs) * 3
+        np.testing.assert_equal(sliced.get_label(), y[1:7])
+        np.testing.assert_equal(sliced.get_float_info('feature_weights'), fw)
+        np.testing.assert_equal(sliced.get_base_margin(), base_margin[1:7, :].flatten())
+        np.testing.assert_equal(sliced.get_base_margin(), sliced.get_float_info('base_margin'))
 
-        eval_res_1 = {}
+        # Slicing a DMatrix results into a DMatrix that's equivalent to a DMatrix that's
-        xgb.train(
+        # constructed from the corresponding NumPy slice
+        d2 = xgb.DMatrix(X[1:7, :], y[1:7])
+        d2.set_base_margin(base_margin[1:7, :].flatten())
+        eval_res = {}
+        _ = xgb.train(
             {'num_class': 3, 'objective': 'multi:softprob',
-             'eval_metric': 'merror'},
+             'eval_metric': 'mlogloss'},
-            sliced,
+            d,
-            num_boost_round=2, evals=[(sliced, 'd')], evals_result=eval_res_1)
+            num_boost_round=2, evals=[(d2, 'd2'), (sliced, 'sliced')], evals_result=eval_res)
-
+        np.testing.assert_equal(eval_res['d2']['mlogloss'], eval_res['sliced']['mlogloss'])
-        eval_res_0 = eval_res_0['d']['merror']
-        eval_res_1 = eval_res_1['d']['merror']
-
-        for i in range(len(eval_res_0)):
-            assert abs(eval_res_0[i] - eval_res_1[i]) < 0.02
 
     def test_feature_names_slice(self):
         data = np.random.randn(5, 5)
 
         # different length
-        self.assertRaises(ValueError, xgb.DMatrix, data,
+        with pytest.raises(ValueError):
-                          feature_names=list('abcdef'))
+            xgb.DMatrix(data, feature_names=list('abcdef'))
         # contains duplicates
-        self.assertRaises(ValueError, xgb.DMatrix, data,
+        with pytest.raises(ValueError):
-                          feature_names=['a', 'b', 'c', 'd', 'd'])
+            xgb.DMatrix(data, feature_names=['a', 'b', 'c', 'd', 'd'])
         # contains symbol
-        self.assertRaises(ValueError, xgb.DMatrix, data,
+        with pytest.raises(ValueError):
-                          feature_names=['a', 'b', 'c', 'd', 'e<1'])
+            xgb.DMatrix(data, feature_names=['a', 'b', 'c', 'd', 'e<1'])
 
         dm = xgb.DMatrix(data)
         dm.feature_names = list('abcde')
@@ -170,14 +167,12 @@ class TestDMatrix(unittest.TestCase):
         dm.feature_types = list('qiqiq')
         assert dm.feature_types == list('qiqiq')
 
-        def incorrect_type_set():
+        with pytest.raises(ValueError):
             dm.feature_types = list('abcde')
 
-        self.assertRaises(ValueError, incorrect_type_set)
-
         # reset
         dm.feature_names = None
-        self.assertEqual(dm.feature_names, ['f0', 'f1', 'f2', 'f3', 'f4'])
+        assert dm.feature_names == ['f0', 'f1', 'f2', 'f3', 'f4']
         assert dm.feature_types is None
 
     def test_feature_names(self):
@@ -209,7 +204,8 @@ class TestDMatrix(unittest.TestCase):
 
             # different feature name must raises error
             dm = xgb.DMatrix(dummy, feature_names=list('abcde'))
-            self.assertRaises(ValueError, bst.predict, dm)
+            with pytest.raises(ValueError):
+                bst.predict(dm)
 
     def test_get_info(self):
         dtrain = xgb.DMatrix(dpath + 'agaricus.txt.train')
@@ -234,9 +230,8 @@ class TestDMatrix(unittest.TestCase):
 
         fw -= 1
 
-        def assign_weight():
+        with pytest.raises(ValueError):
             m.set_info(feature_weights=fw)
-        self.assertRaises(ValueError, assign_weight)
 
     def test_sparse_dmatrix_csr(self):
         nrow = 100

tests/python/test_dt.py

@@ -1,5 +1,4 @@
 # -*- coding: utf-8 -*-
-import unittest
 import pytest
 import numpy as np
 
@@ -17,7 +16,7 @@ pytestmark = pytest.mark.skipif(
     reason=tm.no_dt()['reason'] + ' or ' + tm.no_pandas()['reason'])
 
 
-class TestDataTable(unittest.TestCase):
+class TestDataTable:
 
     def test_dt(self):
         df = pd.DataFrame([[1, 2., True], [2, 3., False]],
@@ -43,7 +42,8 @@ class TestDataTable(unittest.TestCase):
         df = pd.DataFrame([[1, 2., 'x'], [2, 3., 'y']],
                           columns=['a', 'b', 'c'])
         dtable = dt.Frame(df)
-        self.assertRaises(ValueError, xgb.DMatrix, dtable)
+        with pytest.raises(ValueError):
+            xgb.DMatrix(dtable)
 
         df = pd.DataFrame({'A=1': [1, 2, 3], 'A=2': [4, 5, 6]})
         dtable = dt.Frame(df)

tests/python/test_early_stopping.py

@@ -1,13 +1,12 @@
 import xgboost as xgb
 import testing as tm
 import numpy as np
-import unittest
 import pytest
 
 rng = np.random.RandomState(1994)
 
 
-class TestEarlyStopping(unittest.TestCase):
+class TestEarlyStopping:
 
     @pytest.mark.skipif(**tm.no_sklearn())
     def test_early_stopping_nonparallel(self):

tests/python/test_eval_metrics.py

@@ -1,13 +1,12 @@
 import xgboost as xgb
 import testing as tm
 import numpy as np
-import unittest
 import pytest
 
 rng = np.random.RandomState(1337)
 
 
-class TestEvalMetrics(unittest.TestCase):
+class TestEvalMetrics:
     xgb_params_01 = {
         'verbosity': 0,
         'nthread': 1,

tests/python/test_interaction_constraints.py

@@ -1,7 +1,6 @@
 # -*- coding: utf-8 -*-
 import numpy as np
 import xgboost
-import unittest
 import testing as tm
 import pytest
 
@@ -9,7 +8,7 @@ dpath = 'demo/data/'
 rng = np.random.RandomState(1994)
 
 
-class TestInteractionConstraints(unittest.TestCase):
+class TestInteractionConstraints:
     def run_interaction_constraints(self, tree_method):
         x1 = np.random.normal(loc=1.0, scale=1.0, size=1000)
         x2 = np.random.normal(loc=1.0, scale=1.0, size=1000)

tests/python/test_monotone_constraints.py

@@ -1,6 +1,5 @@
 import numpy as np
 import xgboost as xgb
-import unittest
 import testing as tm
 import pytest
 
@@ -61,7 +60,7 @@ y = (
 training_dset = xgb.DMatrix(x, label=y)
 
 
-class TestMonotoneConstraints(unittest.TestCase):
+class TestMonotoneConstraints:
 
     def test_monotone_constraints_for_exact_tree_method(self):
 

tests/python/test_openmp.py

@@ -1,10 +1,9 @@
 # -*- coding: utf-8 -*-
 import xgboost as xgb
-import unittest
 import numpy as np
 
 
-class TestOMP(unittest.TestCase):
+class TestOMP:
     def test_omp(self):
         dpath = 'demo/data/'
         dtrain = xgb.DMatrix(dpath + 'agaricus.txt.train')

tests/python/test_parse_tree.py

@@ -1,5 +1,4 @@
 import xgboost as xgb
-import unittest
 import numpy as np
 import pytest
 import testing as tm
@@ -12,7 +11,7 @@ dpath = 'demo/data/'
 rng = np.random.RandomState(1994)
 
 
-class TestTreesToDataFrame(unittest.TestCase):
+class TestTreesToDataFrame:
 
     def build_model(self, max_depth, num_round):
         dtrain = xgb.DMatrix(dpath + 'agaricus.txt.train')

tests/python/test_pickling.py

@@ -2,7 +2,6 @@ import pickle
 import numpy as np
 import xgboost as xgb
 import os
-import unittest
 
 
 kRows = 100
@@ -15,7 +14,7 @@ def generate_data():
     return X, y
 
 
-class TestPickling(unittest.TestCase):
+class TestPickling:
     def run_model_pickling(self, xgb_params):
         X, y = generate_data()
         dtrain = xgb.DMatrix(X, y)

tests/python/test_plotting.py

@@ -3,7 +3,6 @@ import numpy as np
 import xgboost as xgb
 import testing as tm
 
-import unittest
 import pytest
 
 try:
@@ -20,7 +19,7 @@ pytestmark = pytest.mark.skipif(**tm.no_multiple(tm.no_matplotlib(),
 dpath = 'demo/data/agaricus.txt.train'
 
 
-class TestPlotting(unittest.TestCase):
+class TestPlotting:
     def test_plotting(self):
         m = xgb.DMatrix(dpath)
         booster = xgb.train({'max_depth': 2, 'eta': 1,

tests/python/test_predict.py

@@ -1,5 +1,4 @@
 '''Tests for running inplace prediction.'''
-import unittest
 from concurrent.futures import ThreadPoolExecutor
 import numpy as np
 from scipy import sparse
@@ -66,7 +65,7 @@ def test_predict_leaf():
     run_predict_leaf('cpu_predictor')
 
 
-class TestInplacePredict(unittest.TestCase):
+class TestInplacePredict:
     '''Tests for running inplace prediction'''
     def test_predict(self):
         rows = 1000

tests/python/test_ranking.py

@@ -2,7 +2,6 @@ import numpy as np
 from scipy.sparse import csr_matrix
 import xgboost
 import os
-import unittest
 import itertools
 import shutil
 import urllib.request
@@ -73,10 +72,10 @@ def test_ranking_with_weighted_data():
         assert all(p <= q for p, q in zip(is_sorted, is_sorted[1:]))
 
 
-class TestRanking(unittest.TestCase):
+class TestRanking:
 
     @classmethod
-    def setUpClass(cls):
+    def setup_class(cls):
         """
         Download and setup the test fixtures
         """
@@ -119,7 +118,7 @@ class TestRanking(unittest.TestCase):
                       }
 
     @classmethod
-    def tearDownClass(cls):
+    def teardown_class(cls):
         """
         Cleanup test artifacts from download and unpacking
         :return:
@@ -144,8 +143,9 @@ class TestRanking(unittest.TestCase):
         cv = xgboost.cv(self.params, self.dtrain, num_boost_round=2500,
                         early_stopping_rounds=10, nfold=10, as_pandas=False)
         assert isinstance(cv, dict)
-        self.assertSetEqual(set(cv.keys()), {'test-ndcg-mean', 'train-ndcg-mean', 'test-ndcg-std', 'train-ndcg-std'},
+        assert (set(cv.keys()) == {'test-ndcg-mean', 'train-ndcg-mean', 'test-ndcg-std',
-                            "CV results dict key mismatch")
+                                   'train-ndcg-std'},
+                'CV results dict key mismatch.')
 
     def test_cv_no_shuffle(self):
         """

tests/python/test_shap.py

@@ -1,7 +1,6 @@
 # -*- coding: utf-8 -*-
 import numpy as np
 import xgboost as xgb
-import unittest
 import itertools
 import re
 import scipy
@@ -11,7 +10,7 @@ dpath = 'demo/data/'
 rng = np.random.RandomState(1994)
 
 
-class TestSHAP(unittest.TestCase):
+class TestSHAP:
 
     def test_feature_importances(self):
         data = np.random.randn(100, 5)

tests/python/test_training_continuation.py

@@ -1,13 +1,12 @@
 import xgboost as xgb
 import testing as tm
 import numpy as np
-import unittest
 import pytest
 
 rng = np.random.RandomState(1337)
 
 
-class TestTrainingContinuation(unittest.TestCase):
+class TestTrainingContinuation:
     num_parallel_tree = 3
 
     def generate_parameters(self, use_json):

tests/python/test_tree_regularization.py

@@ -1,5 +1,4 @@
 import numpy as np
-import unittest
 import xgboost as xgb
 
 from numpy.testing import assert_approx_equal
@@ -7,7 +6,7 @@ from numpy.testing import assert_approx_equal
 train_data = xgb.DMatrix(np.array([[1]]), label=np.array([1]))
 
 
-class TestTreeRegularization(unittest.TestCase):
+class TestTreeRegularization:
     def test_alpha(self):
         params = {
             'tree_method': 'exact', 'verbosity': 0,

tests/python/test_updaters.py

@@ -1,5 +1,4 @@
 import testing as tm
-import unittest
 import pytest
 import xgboost as xgb
 import numpy as np
@@ -36,7 +35,7 @@ def train_result(param, dmat, num_rounds):
     return result
 
 
-class TestTreeMethod(unittest.TestCase):
+class TestTreeMethod:
     @given(exact_parameter_strategy, strategies.integers(1, 20),
            tm.dataset_strategy)
     @settings(deadline=None)

tests/python/test_with_modin.py

@@ -2,7 +2,6 @@
 import numpy as np
 import xgboost as xgb
 import testing as tm
-import unittest
 import pytest
 
 try:
@@ -18,7 +17,7 @@ dpath = 'demo/data/'
 rng = np.random.RandomState(1994)
 
 
-class TestModin(unittest.TestCase):
+class TestModin:
 
     def test_modin(self):
 
@@ -43,7 +42,8 @@ class TestModin(unittest.TestCase):
         # incorrect dtypes
         df = md.DataFrame([[1, 2., 'x'], [2, 3., 'y']],
                           columns=['a', 'b', 'c'])
-        self.assertRaises(ValueError, xgb.DMatrix, df)
+        with pytest.raises(ValueError):
+            xgb.DMatrix(df)
 
         # numeric columns
         df = md.DataFrame([[1, 2., True], [2, 3., False]])
@@ -113,13 +113,13 @@ class TestModin(unittest.TestCase):
     def test_modin_label(self):
         # label must be a single column
         df = md.DataFrame({'A': ['X', 'Y', 'Z'], 'B': [1, 2, 3]})
-        self.assertRaises(ValueError, xgb.data._transform_pandas_df, df,
+        with pytest.raises(ValueError):
-                          False, None, None, 'label', 'float')
+            xgb.data._transform_pandas_df(df, False, None, None, 'label', 'float')
 
         # label must be supported dtype
         df = md.DataFrame({'A': np.array(['a', 'b', 'c'], dtype=object)})
-        self.assertRaises(ValueError, xgb.data._transform_pandas_df, df,
+        with pytest.raises(ValueError):
-                          False, None, None, 'label', 'float')
+            xgb.data._transform_pandas_df(df, False, None, None, 'label', 'float')
 
         df = md.DataFrame({'A': np.array([1, 2, 3], dtype=int)})
         result, _, _ = xgb.data._transform_pandas_df(df, False, None, None,

tests/python/test_with_pandas.py

@@ -2,7 +2,6 @@
 import numpy as np
 import xgboost as xgb
 import testing as tm
-import unittest
 import pytest
 
 try:
@@ -18,7 +17,7 @@ dpath = 'demo/data/'
 rng = np.random.RandomState(1994)
 
 
-class TestPandas(unittest.TestCase):
+class TestPandas:
 
     def test_pandas(self):
 
@@ -43,7 +42,8 @@ class TestPandas(unittest.TestCase):
         # incorrect dtypes
         df = pd.DataFrame([[1, 2., 'x'], [2, 3., 'y']],
                           columns=['a', 'b', 'c'])
-        self.assertRaises(ValueError, xgb.DMatrix, df)
+        with pytest.raises(ValueError):
+            xgb.DMatrix(df)
 
         # numeric columns
         df = pd.DataFrame([[1, 2., True], [2, 3., False]])
@@ -139,13 +139,13 @@ class TestPandas(unittest.TestCase):
     def test_pandas_label(self):
         # label must be a single column
         df = pd.DataFrame({'A': ['X', 'Y', 'Z'], 'B': [1, 2, 3]})
-        self.assertRaises(ValueError, xgb.data._transform_pandas_df, df,
+        with pytest.raises(ValueError):
-                          False, None, None, 'label', 'float')
+            xgb.data._transform_pandas_df(df, False, None, None, 'label', 'float')
 
         # label must be supported dtype
         df = pd.DataFrame({'A': np.array(['a', 'b', 'c'], dtype=object)})
-        self.assertRaises(ValueError, xgb.data._transform_pandas_df, df,
+        with pytest.raises(ValueError):
-                          False, None, None, 'label', 'float')
+            xgb.data._transform_pandas_df(df, False, None, None, 'label', 'float')
 
         df = pd.DataFrame({'A': np.array([1, 2, 3], dtype=int)})
         result, _, _ = xgb.data._transform_pandas_df(df, False, None, None,

tests/python/test_with_sklearn.py

@@ -7,7 +7,6 @@ import tempfile
 import os
 import shutil
 import pytest
-import unittest
 import json
 
 rng = np.random.RandomState(1994)
@@ -1012,37 +1011,39 @@ def test_feature_weights():
     assert poly_decreasing[0] < -0.08
 
 
-class TestBoostFromPrediction(unittest.TestCase):
+def run_boost_from_prediction(tree_method):
-    def run_boost_from_prediction(self, tree_method):
+    from sklearn.datasets import load_breast_cancer
-        from sklearn.datasets import load_breast_cancer
+    X, y = load_breast_cancer(return_X_y=True)
-        X, y = load_breast_cancer(return_X_y=True)
+    model_0 = xgb.XGBClassifier(
-        model_0 = xgb.XGBClassifier(
+        learning_rate=0.3, random_state=0, n_estimators=4,
-            learning_rate=0.3, random_state=0, n_estimators=4,
+        tree_method=tree_method)
-            tree_method=tree_method)
+    model_0.fit(X=X, y=y)
-        model_0.fit(X=X, y=y)
+    margin = model_0.predict(X, output_margin=True)
-        margin = model_0.predict(X, output_margin=True)
 
-        model_1 = xgb.XGBClassifier(
+    model_1 = xgb.XGBClassifier(
-            learning_rate=0.3, random_state=0, n_estimators=4,
+        learning_rate=0.3, random_state=0, n_estimators=4,
-            tree_method=tree_method)
+        tree_method=tree_method)
-        model_1.fit(X=X, y=y, base_margin=margin)
+    model_1.fit(X=X, y=y, base_margin=margin)
-        predictions_1 = model_1.predict(X, base_margin=margin)
+    predictions_1 = model_1.predict(X, base_margin=margin)
 
-        cls_2 = xgb.XGBClassifier(
+    cls_2 = xgb.XGBClassifier(
-            learning_rate=0.3, random_state=0, n_estimators=8,
+        learning_rate=0.3, random_state=0, n_estimators=8,
-            tree_method=tree_method)
+        tree_method=tree_method)
-        cls_2.fit(X=X, y=y)
+    cls_2.fit(X=X, y=y)
-        predictions_2 = cls_2.predict(X)
+    predictions_2 = cls_2.predict(X)
-        assert np.all(predictions_1 == predictions_2)
+    assert np.all(predictions_1 == predictions_2)
 
-    @pytest.mark.skipif(**tm.no_sklearn())
-    def test_boost_from_prediction_hist(self):
-        self.run_boost_from_prediction('hist')
 
-    @pytest.mark.skipif(**tm.no_sklearn())
+@pytest.mark.skipif(**tm.no_sklearn())
-    def test_boost_from_prediction_approx(self):
+def test_boost_from_prediction_hist():
-        self.run_boost_from_prediction('approx')
+    run_boost_from_prediction('hist')
 
-    @pytest.mark.skipif(**tm.no_sklearn())
+
-    def test_boost_from_prediction_exact(self):
+@pytest.mark.skipif(**tm.no_sklearn())
-        self.run_boost_from_prediction('exact')
+def test_boost_from_prediction_approx():
+    run_boost_from_prediction('approx')
+
+
+@pytest.mark.skipif(**tm.no_sklearn())
+def test_boost_from_prediction_exact():
+    run_boost_from_prediction('exact')

tests/python/testing.py

@@ -301,6 +301,14 @@ def captured_output():
         sys.stdout, sys.stderr = old_out, old_err
 
 
+try:
+    # Python 3.7+
+    from contextlib import nullcontext as noop_context
+except ImportError:
+    # Python 3.6
+    from contextlib import suppress as noop_context
+
+
 CURDIR = os.path.normpath(os.path.abspath(os.path.dirname(__file__)))
 PROJECT_ROOT = os.path.normpath(
     os.path.join(CURDIR, os.path.pardir, os.path.pardir))