From 31a5954a4a19c877ec74f1805fb1f4a979d78720 Mon Sep 17 00:00:00 2001
From: Phil Elson <philip.elson@cern.ch>
Date: Wed, 5 May 2021 21:03:36 +0200
Subject: [PATCH] Review actions on time handling.

---
 cara/models.py                                | 18 +++----
 cara/tests/models/test_concentration_model.py | 52 +++++++++++--------
 2 files changed, 39 insertions(+), 31 deletions(-)

diff --git a/cara/models.py b/cara/models.py
index 8f2a8c42..9195ef76 100644
--- a/cara/models.py
+++ b/cara/models.py
@@ -644,7 +644,8 @@ class ConcentrationModel:
         k = (vg * 3600) / h
 
         return k + self.virus.decay_constant + self.ventilation.air_exchange(
-                                                    self.room, time)
+            self.room, time
+        )
 
     @cached()
     def _concentration_limit(self, time: float) -> _VectorisedFloat:
@@ -689,7 +690,10 @@ class ConcentrationModel:
         for change_time in self.state_change_times():
             if change_time >= time:
                 return change_time
-        raise ValueError("Time larger than highest state change")
+        raise ValueError(
+            f"The requested time ({time}) is greater than last available "
+            f"state change time ({change_time})"
+        )
 
     @cached()
     def concentration(self, time: float) -> _VectorisedFloat:
@@ -705,13 +709,9 @@ class ConcentrationModel:
 
         if time == 0:
             return 0.0
-        try:
-            IVRR = self.infectious_virus_removal_rate(self._next_state_change(time))
-            concentration_limit = self._concentration_limit(self._next_state_change(time))
-        except ValueError:
-            raise ValueError("Concentration cannot be computed at a time"
-                             " larger than last state change (parameters"
-                             " are not defined)")
+        next_state_change_time = self._next_state_change(time)
+        IVRR = self.infectious_virus_removal_rate(next_state_change_time)
+        concentration_limit = self._concentration_limit(next_state_change_time)
 
         t_last_state_change = self.last_state_change(time)
         concentration_at_last_state_change = self.concentration(t_last_state_change)
diff --git a/cara/tests/models/test_concentration_model.py b/cara/tests/models/test_concentration_model.py
index 0fd81271..be58ddd3 100644
--- a/cara/tests/models/test_concentration_model.py
+++ b/cara/tests/models/test_concentration_model.py
@@ -1,8 +1,6 @@
-from dataclasses import dataclass
 import re
 
 import numpy as np
-import numpy.testing as npt
 import pytest
 
 from cara import models
@@ -62,33 +60,43 @@ def test_concentration_model_vectorisation(override_params):
     assert concentrations.shape == (2, )
 
 
-@pytest.mark.parametrize(
-    "time, expected_next_state_change", [
-        [0, 0],
-        [1, 4],
-        [4, 4],
-        [24, 24],
-    ]
-)
-def test_concentration_model_next_state_change(time,expected_next_state_change):
-    always = models.PeriodicInterval(240, 240)
-    c_model = models.ConcentrationModel(
+@pytest.fixture
+def simple_conc_model():
+    interesting_times = models.SpecificInterval(([0, 1], [1.1, 1.999], [2, 3]), )
+    return models.ConcentrationModel(
         models.Room(75),
-        models.AirChange(always, 100),
+        models.AirChange(interesting_times, 100),
         models.InfectedPopulation(
             number=1,
-            presence=always,
+            presence=interesting_times,
             mask=models.Mask.types['Type I'],
             activity=models.Activity.types['Seated'],
             virus=models.Virus.types['SARS_CoV_2'],
             expiration=models.Expiration.types['Breathing'],
         )
     )
-    assert c_model._next_state_change(time) == expected_next_state_change
-    with pytest.raises(ValueError, match="Time larger than highest state change"):
-        c_model._next_state_change(24.1)
-    with pytest.raises(ValueError, match=re.escape("Concentration cannot "
-            "be computed at a time larger than last state change "
-            "(parameters are not defined)")):
-        c_model.concentration(24.1)
 
+
+@pytest.mark.parametrize(
+    "time, expected_next_state_change", [
+        [0, 0],
+        [1, 1],
+        [1.1, 1.999],
+        [2, 3],
+        [3, 3],
+    ]
+)
+def test_next_state_change_time(
+        simple_conc_model: models.ConcentrationModel,
+        time,
+        expected_next_state_change,
+):
+    simple_conc_model._next_state_change(time) == expected_next_state_change
+
+
+def test_next_state_change_time_out_of_range(simple_conc_model: models.ConcentrationModel):
+    with pytest.raises(
+            ValueError,
+            match=re.escape("The requested time (3.1) is greater than last available state change time (3)")
+    ):
+        simple_conc_model._next_state_change(3.1)