test(model): adjust tests for sir model

#100
kausalflow · Jan 4, 2025 · 442d4d9 · 442d4d9
1 parent 83e6333
commit 442d4d9
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 4 deletions.
diff --git a/hamilflow/models/sir.py b/hamilflow/models/sir.py
@@ -141,7 +141,14 @@ def __call__(self, t: TypeTime) -> pd.DataFrame:
         infected = self.initial_condition.infected_0
         recovered = self.initial_condition.recovered_0
 
-        results = []
+        results = [
+            {
+                "t": 0,
+                "S": susceptible,
+                "I": infected,
+                "R": recovered,
+            },
+        ]
 
         for t_i in np.array(t):
             results.append(

diff --git a/tests/test_models/test_sir.py b/tests/test_models/test_sir.py
@@ -68,13 +68,13 @@ def test_sir_step(sir_model: SIR) -> None:
     delta_s, delta_i, delta_r = sir_model._step(999, 1)
     assert delta_s < 0
     assert delta_i > 0
-    assert delta_r > 0
+    assert delta_r == 0
 
 
 def test_sir_call(sir_model: SIR) -> None:
     """Test the call function of the SIR model."""
-    t = np.arange(0, 10)
+    t = np.arange(1, 10)
     result = sir_model(t)
     assert isinstance(result, pd.DataFrame)
-    assert len(result) == len(t)
+    assert len(result) == len(t) + 1
     assert all(col in result.columns for col in ["t", "S", "I", "R"])