actually passing all the tests

Author: dsweber2

R/arx_forecaster.R

@@ -172,7 +172,7 @@ arx_fcast_epi_workflow <- function(
   r <- r %>%
     step_epi_naomit() %>%
     step_training_window(n_recent = args_list$n_training) %>%
-    check_enough_train_data(all_predictors(), skip = FALSE)
+    check_enough_train_data(all_predictors(), n = args_list$check_enough_data_n, skip = FALSE)
 
   if (!is.null(args_list$check_enough_data_n)) {
     r <- r %>% check_enough_train_data(

R/check_enough_train_data.R

@@ -123,9 +123,11 @@ prep.check_enough_train_data <- function(x, training, info = NULL, ...) {
 bake.check_enough_train_data <- function(object, new_data, ...) {
   col_names <- object$columns
   if (object$drop_na) {
-    newish_data <- tidyr::drop_na(new_data, any_of(unname(col_names)))
+    non_na_data <- tidyr::drop_na(new_data, any_of(unname(col_names)))
+  } else {
+    non_na_data <- new_data
   }
-  cols_not_enough_data <- newish_data %>%
+  cols_not_enough_data <- non_na_data %>%
     group_by(across(all_of(.env$object$epi_keys))) %>%
     summarise(across(all_of(.env$col_names), ~ dplyr::n() < .env$object$n), .groups = "drop") %>%
     summarise(across(all_of(.env$col_names), any), .groups = "drop") %>%

tests/testthat/_snaps/check_enough_train_data.md

@@ -38,7 +38,7 @@
 
     Code
       epi_recipe(toy_epi_df) %>% step_epi_lag(x, lag = c(1, 2)) %>%
-        check_enough_train_data(all_predictors(), y, n = 2 * n - 5) %>% prep(
+        check_enough_train_data(all_predictors(), y, n = 2 * n - 4) %>% prep(
         toy_epi_df) %>% bake(new_data = NULL)
     Condition
       Error in `prep()`:

tests/testthat/test-check_enough_train_data.R

@@ -94,33 +94,34 @@ test_that("check_enough_train_data only checks train data", {
     epiprocess::as_epi_df()
   expect_no_error(
     epi_recipe(toy_epi_df) %>%
-      check_enough_train_data(x, y, n = n - 2, epi_keys = "geo_value") %>%
+      check_enough_train_data(x, y, n = n - 2, epi_keys = "geo_value", skip = TRUE) %>%
       prep(toy_epi_df) %>%
       bake(new_data = toy_test_data)
   )
   # Same thing, but skip = FALSE
   expect_no_error(
     epi_recipe(toy_epi_df) %>%
-      check_enough_train_data(y, n = n - 2, epi_keys = "geo_value", skip = FALSE) %>%
+      check_enough_train_data(y, n = n - 2, epi_keys = "geo_value") %>%
       prep(toy_epi_df) %>%
       bake(new_data = toy_test_data)
   )
 })
 
 test_that("check_enough_train_data works with all_predictors() downstream of constructed terms", {
-  # With a lag of 2, we will get 2 * n - 6 non-NA rows
+  # With a lag of 2, we will get 2 * n - 5 non-NA rows (NA's in x but not in the
+  # lags don't count)
   expect_no_error(
     epi_recipe(toy_epi_df) %>%
       step_epi_lag(x, lag = c(1, 2)) %>%
-      check_enough_train_data(all_predictors(), y, n = 2 * n - 6) %>%
+      check_enough_train_data(all_predictors(), y, n = 2 * n - 5) %>%
       prep(toy_epi_df) %>%
       bake(new_data = NULL)
   )
   expect_snapshot(
     error = TRUE,
     epi_recipe(toy_epi_df) %>%
       step_epi_lag(x, lag = c(1, 2)) %>%
-      check_enough_train_data(all_predictors(), y, n = 2 * n - 5) %>%
+      check_enough_train_data(all_predictors(), y, n = 2 * n - 4) %>%
       prep(toy_epi_df) %>%
       bake(new_data = NULL)
   )

tests/testthat/test-layer_residual_quantiles.R

@@ -103,7 +103,7 @@ test_that("Canned forecasters work with / without", {
   )
 
   expect_silent(
-    arx_forecaster(jhu, "death_rate", c("case_rate", "death_rate"))
+    arx_forecaster(jhu, "death_rate", c("case_rate", "death_rate"), args_list = arx_args_list(check_enough_data_n = 1))
   )
   expect_silent(
     flatline_forecaster(