TST that predictions from predict and multi_predict have same format and all levels

Author: patr1ckm

tests/testthat/test_multinom_reg.R

@@ -104,3 +104,14 @@ test_that('bad input', {
   expect_error(translate(multinom_reg() %>% set_engine()))
   expect_warning(translate(multinom_reg() %>% set_engine("glmnet", x = iris[,1:3], y = iris$Species)))
 })
+
+test_that("predictions are factors with all levels", {
+  basic <- multinom_reg() %>% set_engine("glmnet") %>% fit(Species ~ ., data = iris)
+  nd <- iris[iris$Species == "setosa", ]
+  yhat <- predict(basic, new_data = nd, penalty = .1)
+  expect_is(yhat$.pred_class, "factor")
+  expect_equal(levels(yhat$.pred_class), levels(iris$Species))
+  yhat_multi <- multi_predict(basic, new_data =  nd, penalty = .1)$.pred
+  expect_is(yhat_multi[[1]]$.pred_class, "factor")
+  expect_equal(levels(yhat_multi[[1]]$.pred_class), levels(iris$Species))
+})