Corrected tests.

Added a warning to functions and shorter stop message to make tests work.

R-package/R/lgb.plot.tree.R

@@ -62,8 +62,9 @@ lgb.plot.tree <- function(model = NULL, tree = NULL, rules = NULL) {
     # extract data.table model structure
     modelDT <- lgb.model.dt.tree(model)
     # check that tree is less than or equal to the maximum tree index in the model
-    if (tree > max(modelDT$tree_index)) {
-        stop("lgb.plot.tree: Value of 'tree' should be between 1 and the total number of trees in the model (", max(modelDT$tree_index), "). Got: ", tree, ".")
+    if (tree > max(modelDT$tree_index) || tree < 1) {
+        warning("lgb.plot.tree: Value of 'tree' should be between 1 and the total number of trees in the model (", max(modelDT$tree_index), "). Got: ", tree, ".")
+        stop("lgb.plot.tree: Invalid tree number")
     }
     # filter modelDT to just the rows for the selected tree
     modelDT <- modelDT[tree_index == tree, ]

R-package/tests/testthat/test_lgb.plot.tree.R

@@ -64,24 +64,17 @@ models <- list(
 
 for (model_name in names(models)){
   model <- models[[model_name]]
-  expected_n_trees <- NROUNDS
-  if (model_name == "multi") {
-    expected_n_trees <- NROUNDS * NCLASS
-  }
-  df <- as.data.frame(lgb.model.dt.tree(model))
-  df_list <- split(df, f = df$tree_index, drop = TRUE)
-  df_leaf <- df[!is.na(df$leaf_index), ]
-  df_internal <- df[is.na(df$leaf_index), ]
+  modelDT <- lgb.model.dt.tree(model)
 
   test_that("lgb.plot.tree fails when a non existing tree is selected", {
     expect_error({
       lgb.plot.tree(model, 0)
-    }, regexp = "lgb.plot.tree: Value of 'tree' should be between 1 and the total number of trees in the model")
+    }, regexp = paste0("lgb.plot.tree: Invalid tree number"))
   })
   test_that("lgb.plot.tree fails when a non existing tree is selected", {
     expect_error({
       lgb.plot.tree(model, 999)
-    }, regexp = "lgb.plot.tree: Value of 'tree' should be between 1 and the total number of trees in the model")
+    }, regexp = paste0("lgb.plot.tree: Invalid tree number"))
   })
   test_that("lgb.plot.tree fails when a non numeric tree is selected", {
     expect_error({
@@ -96,7 +89,7 @@ for (model_name in names(models)){
   test_that("lgb.plot.tree fails when a non lgb.Booster model is passed", {
     expect_error({
       lgb.plot.tree(1, 0)
-    }, regexp = "lgb.plot.tree: model should be an 'lgb.Booster'")
+    }, regexp = paste0("lgb.plot.tree: model should be an ", sQuote("lgb.Booster")))
   })
 }