fix the bug @mgyliu caught

Author: dajmcdon

R/epi_keys.R

@@ -29,7 +29,7 @@ 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_chr(mold$extras$roles[molded_names %in% keys], names)
-  unname(mold_keys)
+  mold_keys <- map(mold$extras$roles[molded_names %in% keys], names)
+  unname(unlist(mold_keys))
 }
 

tests/testthat/test-epi_keys.R

@@ -30,3 +30,21 @@ test_that("epi_keys_mold extracts time_value and geo_value, but not raw",{
   expect_equal(epi_keys_mold(my_workflow$pre$mold),
                c("time_value","geo_value"))
 })
+
+test_that("epi_keys_mold extracts additional keys when they are present", {
+  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"))) # <-- metadata constructed here
+
+  my_recipe <- epi_recipe(my_data) %>%
+    step_epi_ahead(death_rate , ahead = 7) %>%
+    step_epi_naomit()
+
+  my_workflow <- epi_workflow(my_recipe, linear_reg()) %>% fit(my_data)
+
+  expect_equal(
+    epi_keys_mold(my_workflow$pre$mold),
+    c("time_value", "geo_value", "state", "pol"))
+})