recovery test

Author: erik-rosenstrom

src/death.rs

@@ -1,54 +1,15 @@
-use ixa::{Context, HashMap, IxaError, plan::PlanId, prelude::*};
+use ixa::{Context, IxaError, prelude::*};
 
 use crate::{
     population_loader::{Alive, Person, PersonId},
     settings::ContextSettingExt,
     symptom_status_manager::SymptomStatus,
 };
 
-#[derive(Default)]
-struct DeathDataContainer {
-    plans: HashMap<PersonId, Vec<PlanId>>,
-}
-
-impl DeathDataContainer {
-    fn record_plan(&mut self, person: PersonId, plan: PlanId) {
-        self.plans.entry(person).or_default().push(plan);
-    }
-
-    fn get_plans(&mut self, person: PersonId) -> Vec<PlanId> {
-        self.plans.get(&person).cloned().unwrap_or_default()
-    }
-
-    fn remove_plan_records(&mut self, person: PersonId) {
-        self.plans.remove(&person);
-    }
-}
-
-define_data_plugin!(
-    DeathDataPlugin,
-    DeathDataContainer,
-    DeathDataContainer::default()
-);
-
 pub trait ContextDeathExt: PluginContext + ContextEntitiesExt {
     fn is_alive(&self, person_id: PersonId) -> bool {
         self.get_property::<Person, Alive>(person_id).0
     }
-
-    fn record_plan(&mut self, person_id: PersonId, plan_id: PlanId) {
-        self.get_data_mut(DeathDataPlugin)
-            .record_plan(person_id, plan_id);
-    }
-
-    fn cancel_plans(&mut self, person_id: PersonId) {
-        let plan_ids = self.get_data_mut(DeathDataPlugin).get_plans(person_id);
-        for plan_id in plan_ids {
-            self.cancel_plan(&plan_id);
-        }
-        self.get_data_mut(DeathDataPlugin)
-            .remove_plan_records(person_id);
-    }
 }
 impl ContextDeathExt for Context {}
 pub fn init(context: &mut Context) -> Result<(), IxaError> {
@@ -214,6 +175,39 @@ mod test {
         context.execute();
     }
 
+    #[test]
+    fn test_dead_person_does_not_recover() {
+        // This test verifies that when a person is dead, they do not transition to a recovered state.
+        // We create a person, set them to dead, and then check that their infection status does not change to recovered over the course of the simulation.
+        let mut context = setup_context(0, 5.0, 0.5, 10.0);
+        let person_id: PersonId = context.add_entity((Age(30),)).unwrap();
+        infection_propagation_loop::init(&mut context).unwrap();
+        symptom_status_manager::init(&mut context).unwrap();
+        init(&mut context).unwrap();
+        // We schedule the person to be set to dead at time 0.0,
+        // and then we check that their infection status does not change to recovered after the death event is processed.
+        context.add_plan_with_phase(
+            0.0,
+            move |context| {
+                assert!(context.is_alive(person_id));
+                context.set_property::<Person, SymptomStatus>(person_id, SymptomStatus::Dead);
+                context.infect_person(person_id, None, None, None);
+            },
+            ExecutionPhase::First,
+        );
+        context.subscribe_to_event::<PropertyChangeEvent<Person, InfectionStatus>>(
+            move |_context, event| {
+                if event.current == InfectionStatus::Recovered {
+                    panic!(
+                        "Dead person should not recover, but person {:?} did",
+                        event.entity_id
+                    );
+                }
+            },
+        );
+        context.execute();
+    }
+
     #[test]
     fn test_no_infection_when_dead() {
         // This test verifies that when a person is dead, they cannot infect others.
@@ -249,11 +243,11 @@ mod test {
             // Add a watcher if the other person is infected.
             context.subscribe_to_event::<PropertyChangeEvent<Person, InfectionStatus>>(
                 move |context, event| {
-                    if event.current == InfectionStatus::Infectious {
-                        if event.entity_id != infectious_person {
-                            *num_infected_clone.borrow_mut() += 1;
-                            context.set_property(event.entity_id, InfectionData::Susceptible);
-                        }
+                    if event.current == InfectionStatus::Infectious
+                        && event.entity_id != infectious_person
+                    {
+                        *num_infected_clone.borrow_mut() += 1;
+                        context.set_property(event.entity_id, InfectionData::Susceptible);
                     }
                 },
             );

src/infection_propagation_loop.rs

@@ -19,6 +19,9 @@ fn schedule_next_forecasted_infection(context: &mut Context, person: PersonId) {
     }) = get_forecast(context, person)
     {
         context.add_plan(next_time, move |context| {
+            if !context.is_alive(person) {
+                return;
+            }
             let _span = open_span("evaluate and schedule next forecast");
             if evaluate_forecast(context, person, forecasted_total_infectiousness) {
                 let _ = infection_attempt(context, person);