diff --git a/caimira/tests/test_conditional_probability.py b/caimira/tests/test_conditional_probability.py
index f9da0c9a..cae5f057 100644
--- a/caimira/tests/test_conditional_probability.py
+++ b/caimira/tests/test_conditional_probability.py
@@ -45,14 +45,14 @@ def baseline_exposure_model():
 def test_conditional_prob_inf_given_vl_dist(baseline_exposure_model):
     
     viral_loads = np.array([3., 5., 7., 9.,])
-    mc_model = baseline_exposure_model.build_model(250_000)
+    mc_model: models.ExposureModel = baseline_exposure_model.build_model(2_000_000)
 
     expected_pi_means = []
     expected_lower_percentiles = []
     expected_upper_percentiles = []
 
     for vl in viral_loads:
-        model_vl: mc.ExposureModel = nested_replace(
+        model_vl: models.ExposureModel = nested_replace(
             mc_model, {
                 'concentration_model.infected.virus.viral_load_in_sputum' : 10**vl,
             }