added unit tests

Author: topepo

tests/testthat/test_nearest_neighbor_kknn.R

@@ -113,3 +113,78 @@ test_that('kknn prediction', {
 
   expect_equal(form_pred, predict(res_form, iris[1:5, c("Sepal.Width", "Species")])$.pred)
 })
+
+
+test_that('kknn multi-predict', {
+
+  skip_if_not_installed("kknn")
+  library(kknn)
+
+  iris_te <- c(1:2, 50:51, 100:101)
+  k_vals <- 1:10
+
+  res_xy <- fit_xy(
+    nearest_neighbor(mode = "classification", neighbors = 3) %>%
+      set_engine("kknn"),
+    control = ctrl,
+    x = iris[-iris_te, num_pred],
+    y = iris$Species[-iris_te]
+  )
+
+  pred_multi <- multi_predict(res_xy, iris[iris_te, num_pred], neighbors = k_vals)
+  expect_equal(pred_multi %>% unnest() %>% nrow(), length(iris_te) * length(k_vals))
+  expect_equal(pred_multi %>% nrow(), length(iris_te))
+
+  pred_uni <- predict(res_xy, iris[iris_te, num_pred])
+  pred_uni_obs <-
+    pred_multi %>%
+    mutate(.rows = row_number()) %>%
+    unnest() %>%
+    dplyr::filter(neighbors == 3) %>%
+    arrange(.rows) %>%
+    dplyr::select(.pred_class)
+  expect_equal(pred_uni, pred_uni_obs)
+
+
+  prob_multi <- multi_predict(res_xy, iris[iris_te, num_pred],
+                              neighbors = k_vals, type = "prob")
+  expect_equal(prob_multi %>% unnest() %>% nrow(), length(iris_te) * length(k_vals))
+  expect_equal(prob_multi %>% nrow(), length(iris_te))
+
+  prob_uni <- predict(res_xy, iris[iris_te, num_pred], type = "prob")
+  prob_uni_obs <-
+    prob_multi %>%
+    mutate(.rows = row_number()) %>%
+    unnest() %>%
+    dplyr::filter(neighbors == 3) %>%
+    arrange(.rows) %>%
+    dplyr::select(!!names(prob_uni))
+  expect_equal(prob_uni, prob_uni_obs)
+
+  # ----------------------------------------------------------------------------
+  # regression
+
+  cars_te <- 1:5
+  k_vals <- 1:10
+
+  res_xy <- fit(
+    nearest_neighbor(mode = "regression", neighbors = 3) %>%
+      set_engine("kknn"),
+    control = ctrl,
+    mpg ~ ., data = mtcars[-cars_te, ]
+  )
+
+  pred_multi <- multi_predict(res_xy, mtcars[cars_te, -1], neighbors = k_vals)
+  expect_equal(pred_multi %>% unnest() %>% nrow(), length(cars_te) * length(k_vals))
+  expect_equal(pred_multi %>% nrow(), length(cars_te))
+
+  pred_uni <- predict(res_xy, mtcars[cars_te, -1])
+  pred_uni_obs <-
+    pred_multi %>%
+    mutate(.rows = row_number()) %>%
+    unnest() %>%
+    dplyr::filter(neighbors == 3) %>%
+    arrange(.rows) %>%
+    dplyr::select(.pred)
+  expect_equal(pred_uni, pred_uni_obs)
+})