Make test single initialization

tests/python_package_test/test_engine.py

@@ -24,7 +24,6 @@
 
 from .utils import (
     SERIALIZERS,
-    builtin_objective,
     dummy_obj,
     load_breast_cancer,
     load_digits,
@@ -4411,6 +4410,20 @@ def test_quantized_training():
 )
 @pytest.mark.skipif(getenv("TASK", "") == "cuda", reason="Skip due to ObjectiveFunction not exposed for cuda devices.")
 def test_objective_function_class(use_weight, num_boost_round, custom_objective, objective_name, df, num_class):
+    def builtin_objective(name, params):
+        fobj = lgb.ObjectiveFunction(name, params)
+
+        def loss(y_pred, dtrain):
+            if fobj.num_data is None:
+                fobj.init(dtrain)
+            (grad, hess) = fobj.get_gradients(y_pred)
+            if fobj.num_class != 1:
+                grad = grad.reshape((fobj.num_class, -1)).transpose()
+                hess = hess.reshape((fobj.num_class, -1)).transpose()
+            return (grad, hess)
+
+        return loss
+
     X, y = df
     rng = np.random.default_rng()
     weight = rng.choice([1, 2], y.shape) if use_weight else None

tests/python_package_test/utils.py

@@ -168,21 +168,6 @@ def multiclass_custom_objective(y_pred, ds):
     return grad, hess
 
 
-def builtin_objective(name, params):
-    """Mimics the builtin objective functions to mock training."""
-
-    def wrapper(y_pred, dtrain):
-        fobj = lgb.ObjectiveFunction(name, params)
-        fobj.init(dtrain)
-        (grad, hess) = fobj.get_gradients(y_pred)
-        if fobj.num_class != 1:
-            grad = grad.reshape((fobj.num_class, -1)).transpose()
-            hess = hess.reshape((fobj.num_class, -1)).transpose()
-        return (grad, hess)
-
-    return wrapper
-
-
 def pickle_obj(obj, filepath, serializer):
     if serializer == "pickle":
         with open(filepath, "wb") as f: