Adding second coffee break in test on exposed presence interval (test_model_generator.py)

cara/apps/calculator/model_generator.py

@@ -174,7 +174,7 @@ class FormData:
             coffee_times.append((start, end))
         return tuple(coffee_times)
 
-    def present_interval(self) -> models.Interval:
+    def infected_present_interval(self) -> models.Interval:
         leave_times = []
         enter_times = []
         if self.lunch_option:
@@ -223,6 +223,55 @@ class FormData:
 
         return models.SpecificInterval(tuple(present_intervals))
 
+    def exposed_present_interval(self) -> models.Interval:
+        leave_times = []
+        enter_times = []
+        if self.lunch_option:
+            leave_times.append(self.lunch_start)
+            enter_times.append(self.lunch_finish)
+
+        for coffee_start, coffee_end in self.coffee_break_times():
+            leave_times.append(coffee_start)
+            enter_times.append(coffee_end)
+
+        # These lists represent the times where the infected person leaves or enters the room, respectively, sorted in
+        # reverse order. Note that these lists allows the person to "leave" when they should not even be present in the
+        # room. The following loop handles this.
+        leave_times.sort(reverse=True)
+        enter_times.sort(reverse=True)
+
+        # This loop iterates through the lists above, populating present_intervals with (enter, leave) intervals
+        # representing the infected person entering and leaving the room. Note that if one of the evenly spaced coffee-
+        # breaks happens to coincide with the lunch-break, it is simply ignored.
+        present_intervals = []
+        time = self.activity_start
+        is_present = True
+        while time < self.activity_finish:
+            if is_present:
+                if not leave_times:
+                    present_intervals.append((time / 60, self.activity_finish / 60))
+                    break
+
+                if leave_times[-1] <= time:
+                    leave_times.pop()
+                else:
+                    new_time = leave_times.pop()
+                    present_intervals.append((time / 60, min(new_time, self.activity_finish) / 60))
+                    is_present = False
+                    time = new_time
+
+            else:
+                if not enter_times:
+                    break
+
+                if enter_times[-1] < time:
+                    enter_times.pop()
+                else:
+                    is_present = True
+                    time = enter_times.pop()
+
+        return models.SpecificInterval(tuple(present_intervals))
+
 
 def model_from_form(form: FormData) -> models.ExposureModel:
     # Initializes room with volume either given directly or as product of area and height
@@ -264,7 +313,7 @@ def model_from_form(form: FormData) -> models.ExposureModel:
             infected=models.InfectedPopulation(
                 number=infected_occupants,
                 virus=virus,
-                presence=form.present_interval(),
+                presence=form.infected_present_interval(),
                 mask=mask,
                 activity=infected_activity,
                 expiration=infected_expiration
@@ -272,7 +321,7 @@ def model_from_form(form: FormData) -> models.ExposureModel:
         ),
         exposed=models.Population(
             number=exposed_occupants,
-            presence=form.present_interval(),
+            presence=form.exposed_present_interval(),
             activity=exposed_activity,
             mask=mask,
         )

cara/tests/apps/calculator/test_model_generator.py

@@ -120,7 +120,7 @@ def test_exposed_present_intervals(baseline_form):
     baseline_form.lunch_finish = 13 * 60 + 30
     baseline_form.infected_start = 10 * 60
     baseline_form.infected_finish = 15 * 60
-    correct = ((9, 11), (11.25, 12.5), (13.5, 17.0))
+    correct = ((9, 11), (11.25, 12.5), (13.5, 15), (15.25, 17.0))
     assert baseline_form.exposed_present_interval().present_times == correct
 
 def test_key_validation(baseline_form_data):