Pull in updated unit test

pipelines/tests/test_build_pyrenew_model.py

@@ -16,6 +16,9 @@ def mock_data():
             "state_pop": [7, 8, 9],
             "generation_interval_pmf": [0.1, 0.2, 0.7],
             "inf_to_ed_pmf": [0.4, 0.5, 0.1],
+            "inf_to_hosp_admit_pmf": [0.0, 0.7, 0.1, 0.1, 0.1],
+            "inf_to_hosp_admit_lognormal_loc": 0.015,
+            "inf_to_hosp_admit_lognormal_scale": 0.851,
             "right_truncation_pmf": [0.7, 0.1, 0.2],
             "nssp_training_dates": ["2025-01-01"],
             "nhsn_training_dates": ["2025-01-02"],
@@ -44,6 +47,8 @@ def mock_priors():
 ihr_rv = None
 t_peak_rv = None
 duration_shed_after_peak_rv = None
+inf_to_ed_offset_loc_rv = None
+inf_to_ed_log_offset_scale_rv = None
 log10_genome_per_inf_ind_rv = None
 mode_sigma_ww_site_rv = None
 sd_log_sigma_ww_site_rv = None
@@ -65,18 +70,18 @@ def test_build_model_from_dir(tmp_path, mock_data, mock_priors):
 
     model_data = json.loads(mock_data)
 
-    # Test when all sample arguments are False
+    # Test when all `fit_` arguments are False
     _, data = build_model_from_dir(model_dir)
     assert data.data_observed_disease_ed_visits is None
     assert data.data_observed_disease_hospital_admissions is None
     assert data.data_observed_disease_wastewater is None
 
-    # Test when all sample arguments are True
+    # Test when all `fit_` arguments are True
     _, data = build_model_from_dir(
         model_dir,
-        sample_ed_visits=True,
-        sample_hospital_admissions=True,
-        sample_wastewater=True,
+        fit_ed_visits=True,
+        fit_hospital_admissions=True,
+        fit_wastewater=True,
     )
     assert jnp.array_equal(
         data.data_observed_disease_ed_visits,