From 236a7139f54ef1651116791ab5398eccb10f9a02 Mon Sep 17 00:00:00 2001
From: Nicolas Mounet <nicolas.mounet@cern.ch>
Date: Tue, 4 May 2021 08:49:02 +0200
Subject: [PATCH] Adding a test for next_state_change in ConcentrationModel

---
 cara/models.py                                |  2 +-
 cara/tests/models/test_concentration_model.py | 27 +++++++++++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/cara/models.py b/cara/models.py
index 553dd8bd..609da9a3 100644
--- a/cara/models.py
+++ b/cara/models.py
@@ -684,7 +684,7 @@ class ConcentrationModel:
         for change_time in self.state_change_times():
             if change_time >= time:
                 return change_time
-        raise ValueError("Time higher than largest state change")
+        raise ValueError("Time larger than highest state change")
 
     @cached()
     def concentration(self, time: float) -> _VectorisedFloat:
diff --git a/cara/tests/models/test_concentration_model.py b/cara/tests/models/test_concentration_model.py
index 90bc86e4..10cf6a8c 100644
--- a/cara/tests/models/test_concentration_model.py
+++ b/cara/tests/models/test_concentration_model.py
@@ -61,6 +61,33 @@ 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(
+        models.Room(75),
+        models.AirChange(always, 100),
+        models.InfectedPopulation(
+            number=1,
+            presence=always,
+            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)
+
+
 @dataclass(frozen=True)
 class DummyVentilation(models.Ventilation):
     # Dummy ventilation where air_exchange depends on time explicitly