change test location and add additional shape test

tests/python_package_test/test_engine.py

@@ -15,7 +15,7 @@
 import psutil
 import pytest
 from scipy.sparse import csr_matrix, isspmatrix_csc, isspmatrix_csr
-from sklearn.datasets import load_svmlight_file, make_blobs, make_multilabel_classification, make_regression
+from sklearn.datasets import load_svmlight_file, make_blobs, make_classification, make_multilabel_classification, make_regression
 from sklearn.metrics import average_precision_score, log_loss, mean_absolute_error, mean_squared_error, roc_auc_score
 from sklearn.model_selection import GroupKFold, TimeSeriesSplit, train_test_split
 
@@ -2334,14 +2334,6 @@ def test_refit_with_one_tree_multiclass_classification():
     assert isinstance(model_refit, lgb.Booster)
 
 
-def test_pred_leaf_output_shape():
-    X, y = make_regression(n_samples=10_000, n_features=10)
-    dtrain = lgb.Dataset(X, label=y)
-    params = {"objective": "regression", "verbosity": -1}
-    assert lgb.train(params, dtrain, num_boost_round=1).predict(X, pred_leaf=True).shape == (10_000, 1)
-    assert lgb.train(params, dtrain, num_boost_round=2).predict(X, pred_leaf=True).shape == (10_000, 2)
-
-
 def test_refit_dataset_params(rng):
     # check refit accepts dataset_params
     X, y = load_breast_cancer(return_X_y=True)
@@ -3886,6 +3878,69 @@ def test_predict_stump(rng, use_init_score):
         np.testing.assert_allclose(preds_all, np.full_like(preds_all, fill_value=y_avg))
 
 
+def test_predict_regression_output_shape():
+    X, y = make_regression(n_samples=10_000, n_features=10)
+    dtrain = lgb.Dataset(X, label=y)
+    params = {"objective": "regression", "verbosity": -1}
+    assert lgb.train(params, dtrain, num_boost_round=1).predict(X).shape == (10_000,)
+
+
+def test_predict_binary_classification_output_shape():
+    X, y = make_classification(n_samples=10_000, n_features=10, n_classes=2)
+    dtrain = lgb.Dataset(X, label=y)
+    params = {"objective": "binary", "verbosity": -1}
+    assert lgb.train(params, dtrain, num_boost_round=1).predict(X).shape == (10_000,)
+
+
+def test_predict_multiclass_classification_output_shape():
+    X, y = make_classification(n_samples=10_000, n_features=10, n_classes=3, n_informative=6)
+    dtrain = lgb.Dataset(X, label=y)
+    params = {"objective": "multiclass", "verbosity": -1, "num_class": 3}
+    assert lgb.train(params, dtrain, num_boost_round=1).predict(X).shape == (10_000, 3)
+
+
+def test_predict_leaf_regression_output_shape():
+    X, y = make_regression(n_samples=10_000, n_features=10)
+    dtrain = lgb.Dataset(X, label=y)
+    params = {"objective": "regression", "verbosity": -1}
+    assert lgb.train(params, dtrain, num_boost_round=1).predict(X, pred_leaf=True).shape == (10_000, 1)
+
+
+def test_predict_leaf_binary_classification_output_shape():
+    X, y = make_classification(n_samples=10_000, n_features=10, n_classes=2)
+    dtrain = lgb.Dataset(X, label=y)
+    params = {"objective": "binary", "verbosity": -1}
+    assert lgb.train(params, dtrain, num_boost_round=1).predict(X, pred_leaf=True).shape == (10_000, 1)
+
+
+def test_predict_leaf_multiclass_classification_output_shape():
+    X, y = make_classification(n_samples=10_000, n_features=10, n_classes=3, n_informative=6)
+    dtrain = lgb.Dataset(X, label=y)
+    params = {"objective": "multiclass", "verbosity": -1, "num_class": 3}
+    assert lgb.train(params, dtrain, num_boost_round=1).predict(X, pred_leaf=True).shape == (10_000, 3)
+
+
+def test_predict_contrib_regression_output_shape():
+    X, y = make_regression(n_samples=10_000, n_features=10)
+    dtrain = lgb.Dataset(X, label=y)
+    params = {"objective": "regression", "verbosity": -1}
+    assert lgb.train(params, dtrain, num_boost_round=1).predict(X, pred_contrib=True).shape == (10_000, 11)
+
+
+def test_predict_contrib_binary_classification_output_shape():
+    X, y = make_classification(n_samples=10_000, n_features=10, n_classes=2)
+    dtrain = lgb.Dataset(X, label=y)
+    params = {"objective": "binary", "verbosity": -1}
+    assert lgb.train(params, dtrain, num_boost_round=1).predict(X, pred_contrib=True).shape == (10_000, 11)
+
+
+def test_predict_contrib_multiclass_classification_output_shape():
+    X, y = make_classification(n_samples=10_000, n_features=10, n_classes=3, n_informative=6)
+    dtrain = lgb.Dataset(X, label=y)
+    params = {"objective": "multiclass", "verbosity": -1, "num_class": 3}
+    assert lgb.train(params, dtrain, num_boost_round=1).predict(X, pred_contrib=True).shape == (10_000, 33)
+
+
 def test_average_precision_metric():
     # test against sklearn average precision metric
     X, y = load_breast_cancer(return_X_y=True)