Working predict hooks and tests.

syberia · Aug 2, 2015 · 1adb4ab · 1adb4ab
1 parent 8a175b0
commit 1adb4ab
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 2 deletions.
diff --git a/R/tundraContainer-predict.R b/R/tundraContainer-predict.R
@@ -29,7 +29,7 @@ predict <- function(dataframe, predict_args = list(), verbose = FALSE, munge = T
 
   private$run_hooks("predict_pre_munge")
   if (isTRUE(munge) && length(self$.munge_procedure) > 0) {
-    initial_nrow <- NROW(datafram)
+    initial_nrow <- NROW(dataframe)
     dataframe <- munge(dataframe, self$.munge_procedure, verbose)
     if (NROW(dataframe) != initial_nrow) {
       warning("Some rows were removed during data preparation. ",
@@ -38,7 +38,7 @@ predict <- function(dataframe, predict_args = list(), verbose = FALSE, munge = T
   }
   private$run_hooks("predict_post_munge")
 
-  if (length(formals(self$.predict_function) < 2 || missing(predict_args))) {
+  if (length(formals(self$.predict_function)) < 2 || missing(predict_args)) {
     args <- list(dataframe)
   } else {
     args <- list(dataframe, predict_args)

diff --git a/tests/testthat/test-hooks.R b/tests/testthat/test-hooks.R
@@ -61,3 +61,25 @@ test_that("it runs the train_finalize hook in the correct order", {
   expect_equal(env$effect, c("train", "hook"))
 })
 
+test_that("it runs the predict pre-munge hook in the correct order", {
+  env  <- list2env(list(effect = character(0)))
+  hook <- function() { env$effect <- c(env$effect, "hook") }
+  container <- container_with_munge_side_effect(env)
+  container$add_hook("predict_pre_munge", hook)
+  container$train(iris, munge = FALSE)
+  expect_equal(env$effect, character(0))
+  container$predict(iris)
+  expect_equal(env$effect, c("hook", "mungebit"))
+})
+
+test_that("it runs the predict post-munge hook in the correct order", {
+  env  <- list2env(list(effect = character(0)))
+  hook <- function() { env$effect <- c(env$effect, "hook") }
+  container <- container_with_munge_side_effect(env)
+  container$add_hook("predict_post_munge", hook)
+  container$train(iris, munge = FALSE)
+  expect_equal(env$effect, character(0))
+  container$predict(iris)
+  expect_equal(env$effect, c("mungebit", "hook"))
+})
+