Update tests, adding tests that would have caught previous bug

hewr/tests/testthat/test_to_epiweekly_quantile_table.R

@@ -39,9 +39,20 @@ test_that("to_epiweekly_quantiles works as expected", {
     ) |> suppressMessages()
 
     expect_s3_class(result, "tbl_df")
-    expect_setequal(c(
-      "epiweek", "epiyear", "quantile_value", "quantile_level", "location"
-    ), colnames(result))
+    checkmate::expect_names(
+      colnames(result),
+      identical.to = c(
+        "epiweek",
+        "epiyear",
+        "quantile_value",
+        "quantile_level",
+        "location",
+        "source_samples"
+      )
+    )
+
+    expect_equal(draws_file_name, unique(result$source_samples))
+
     expect_gt(nrow(result), 0)
   }
 
@@ -127,7 +138,11 @@ test_that("to_epiweekly_quantiles handles missing forecast files", {
 
 
 # tests for `to_epiweekly_quantile_table`
-test_that("to_epiweekly_quantile_table handles multiple locations", {
+test_that(paste0(
+  "to_epiweekly_quantile_table ",
+  "handles multiple locations ",
+  "and multiple source files"
+), {
   batch_dir_name <- "covid-19_r_2024-12-14_f_2024-12-08_t_2024-12-14"
   tempdir <- withr::local_tempdir()
 
@@ -142,6 +157,17 @@ test_that("to_epiweekly_quantile_table handles multiple locations", {
     if (loc != "loc3") {
       disease_cols <- c(disease_cols, "prop_disease_ed_visits")
     }
+    create_tidy_forecast_data(
+      directory = loc_dir,
+      filename = "epiweekly_with_epiweekly_other_samples.parquet",
+      date_cols = seq(
+        lubridate::ymd("2024-12-08"), lubridate::ymd("2024-12-14"),
+        by = "week"
+      ),
+      disease_cols = disease_cols,
+      n_draw = 25,
+      with_epiweek = TRUE
+    )
 
     create_tidy_forecast_data(
       directory = loc_dir,
@@ -157,7 +183,10 @@ test_that("to_epiweekly_quantile_table handles multiple locations", {
   })
 
   ## should succeed despite loc3 not having valid draws with strict = FALSE
-  result_w_both_locations <- to_epiweekly_quantile_table(temp_batch_dir) |>
+  result_w_both_locations <-
+    to_epiweekly_quantile_table(temp_batch_dir,
+      epiweekly_other_locations = "loc1"
+    ) |>
     suppressMessages()
 
   ## should error if strict = TRUE because loc3 does not have
@@ -168,6 +197,36 @@ test_that("to_epiweekly_quantile_table handles multiple locations", {
     "did not find valid draws"
   )
 
+  ## should succeed with strict = TRUE if loc3 is excluded
+  alt_result_w_both_locations <- (
+    to_epiweekly_quantile_table(temp_batch_dir,
+      strict = TRUE,
+      exclude = "loc3"
+  )) |>
+    suppressMessages()
+
+  ## results should be equivalent for loc2,
+  ## but not for loc1
+  expect_equal(
+    result_w_both_locations |>
+      dplyr::filter(location == "loc2"),
+    alt_result_w_both_locations |>
+      dplyr::filter(location == "loc2")
+  )
+
+  ## check that one used epiweekly
+  ## other for loc1 while other used
+  ## default, resulting in different values
+  loc1_a <- result_w_both_locations |>
+    dplyr::filter(location == "loc1") |>
+    dplyr::pull(.data$value)
+  loc1_b <- alt_result_w_both_locations |>
+    dplyr::filter(location == "loc1") |>
+    dplyr::pull(.data$value)
+  expect_false(
+    any(loc1_a == loc1_b)
+  )
+
   expect_s3_class(result_w_both_locations, "tbl_df")
   expect_gt(nrow(result_w_both_locations), 0)
   checkmate::expect_names(
@@ -185,16 +244,28 @@ test_that("to_epiweekly_quantile_table handles multiple locations", {
     )
   )
   expect_setequal(
-    c("loc1", "loc2"),
-    result_w_both_locations$location
+    result_w_both_locations$location,
+    c("loc1", "loc2")
+  )
+  expect_setequal(
+    alt_result_w_both_locations$location,
+    c("loc1", "loc2")
   )
-  expect_false("loc3" %in% result_w_both_locations$location)
 
-  result_w_one_location <- to_epiweekly_quantile_table(
-    model_batch_dir = temp_batch_dir,
-    exclude = "loc1"
-  ) |>
-    suppressMessages()
-  expect_true("loc2" %in% result_w_one_location$location)
-  expect_false("loc1" %in% result_w_one_location$location)
+  expect_setequal(
+    result_w_both_locations$source_samples,
+    c(
+      "epiweekly_samples",
+      "epiweekly_with_epiweekly_other_samples"
+    )
+  )
+
+  expect_setequal(
+    alt_result_w_both_locations$source_samples,
+    "epiweekly_samples"
+  )
+
+
+  expect_false("loc3" %in% result_w_both_locations$location)
+  expect_false("loc3" %in% alt_result_w_both_locations$location)
 })