update epi_keys_mold and tests to handle additional keys

Author: mgyliu

R/epi_keys.R

@@ -29,7 +29,8 @@ epi_keys.recipe <- function(x) {
 epi_keys_mold <- function(mold) {
   keys <- c("time_value", "geo_value", "key")
   molded_names <- names(mold$extras$roles)
+  
   mold_keys <- map(mold$extras$roles[molded_names %in% keys], names)
-  unname(unlist(mold_keys))
+  unlist(unname(mold_keys))
 }
 

tests/testthat/test-epi_keys.R

@@ -32,23 +32,19 @@ test_that("epi_keys_mold extracts time_value and geo_value, but not raw",{
 })
 
 test_that("epi_keys_mold extracts additional keys when they are present", {
-  my_data <- tibble::tibble(
-    geo_value = rep(c("ca", "fl", "pa"), each = 3),
-    time_value = rep(seq(as.Date("2020-06-01"), as.Date("2020-06-03"),
-                         by = "day"), length.out = length(geo_value)),
-    pol = rep(c("blue", "swing", "swing"), each = 3), # extra key
-    state = rep(c("ca", "fl", "pa"), each = 3), # extra key
-    value = 1:length(geo_value) + 0.01 * rnorm(length(geo_value))
-  ) %>%
-    epiprocess::as_epi_df(additional_metadata = list(other_keys = c("state", "pol")))
-
+  my_data <- case_death_rate_subset %>%
+    sample_n(6) %>%
+    tsibble::as_tsibble() %>% # add 2 extra keys
+    mutate(state = rep("MA", 6), pol = rep("blue", 6)) %>% 
+    as_epi_df(additional_metadata = list(other_keys=c("state", "pol")))
+  
   my_recipe <- epi_recipe(my_data) %>%
-    step_epi_ahead(value , ahead = 7) %>%
+    step_epi_ahead(death_rate , ahead = 7) %>%
     step_epi_naomit()
 
   my_workflow <- epi_workflow(my_recipe, linear_reg()) %>% fit(my_data)
 
-  expect_setequal(
-    epi_keys_mold(my_workflow$pre$mold),
+  expect_equal(
+    epi_keys_mold(my_workflow$pre$mold), 
     c("time_value", "geo_value", "state", "pol"))
 })