From bf37d783bc19def4cd6d1216798deba8b3b104e3 Mon Sep 17 00:00:00 2001
From: Nikita Titov <nekit94-08@mail.ru>
Date: Sun, 22 Dec 2024 19:44:39 +0300
Subject: [PATCH] Update test_engine.py

---
 tests/python_package_test/test_engine.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/tests/python_package_test/test_engine.py b/tests/python_package_test/test_engine.py
index 667cb86c1a14..47d1f5e3349d 100644
--- a/tests/python_package_test/test_engine.py
+++ b/tests/python_package_test/test_engine.py
@@ -3909,12 +3909,14 @@ def test_predict_regression_output_shape():
     # 1-round model
     bst = lgb.train(params, dtrain, num_boost_round=1)
     assert bst.predict(X).shape == (n_samples,)
+    assert bst.predict(X, raw_score=True).shape == (n_samples,)
     assert bst.predict(X, pred_contrib=True).shape == (n_samples, n_features + 1)
     assert bst.predict(X, pred_leaf=True).shape == (n_samples, 1)
 
     # 2-round model
     bst = lgb.train(params, dtrain, num_boost_round=2)
     assert bst.predict(X).shape == (n_samples,)
+    assert bst.predict(X, raw_score=True).shape == (n_samples,)
     assert bst.predict(X, pred_contrib=True).shape == (n_samples, n_features + 1)
     assert bst.predict(X, pred_leaf=True).shape == (n_samples, 2)
 
@@ -3929,12 +3931,14 @@ def test_predict_binary_classification_output_shape():
     # 1-round model
     bst = lgb.train(params, dtrain, num_boost_round=1)
     assert bst.predict(X).shape == (n_samples,)
+    assert bst.predict(X, raw_score=True).shape == (n_samples,)
     assert bst.predict(X, pred_contrib=True).shape == (n_samples, n_features + 1)
     assert bst.predict(X, pred_leaf=True).shape == (n_samples, 1)
 
     # 2-round model
     bst = lgb.train(params, dtrain, num_boost_round=2)
     assert bst.predict(X).shape == (n_samples,)
+    assert bst.predict(X, raw_score=True).shape == (n_samples,)
     assert bst.predict(X, pred_contrib=True).shape == (n_samples, n_features + 1)
     assert bst.predict(X, pred_leaf=True).shape == (n_samples, 2)
 
@@ -3950,12 +3954,14 @@ def test_predict_multiclass_classification_output_shape():
     # 1-round model
     bst = lgb.train(params, dtrain, num_boost_round=1)
     assert bst.predict(X).shape == (n_samples, n_classes)
+    assert bst.predict(X, raw_score=True).shape == (n_samples, n_classes)
     assert bst.predict(X, pred_contrib=True).shape == (n_samples, n_classes * (n_features + 1))
     assert bst.predict(X, pred_leaf=True).shape == (n_samples, n_classes)
 
     # 2-round model
     bst = lgb.train(params, dtrain, num_boost_round=2)
     assert bst.predict(X).shape == (n_samples, n_classes)
+    assert bst.predict(X, raw_score=True).shape == (n_samples, n_classes)
     assert bst.predict(X, pred_contrib=True).shape == (n_samples, n_classes * (n_features + 1))
     assert bst.predict(X, pred_leaf=True).shape == (n_samples, n_classes * 2)