Merge pull request #51 from djgroen/testing

Testing: perriodic update

facs/base/household.py

@@ -5,7 +5,8 @@
 import random
 
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
+from warnings import warn
 
 from .person import Person
 from .utils import probability
@@ -30,29 +31,34 @@ class Household:
 
     house: House
     ages: list[float]
-    size: int = -1
+    size: Optional[int] = None
 
     agents: list[Person] = field(default_factory=list, init=False)
 
     def __post_init__(self):
         """Post init function."""
 
-        if self.size <= -1:
+        if self.size is None:
             self.size = random.choice([1, 2, 3, 4])
 
+        if self.size < 1:
+            raise ValueError("Household size must be at least 1.")
+
+        if self.size > 4:
+            warn("Household size is greater than 4.")
+
         for _ in range(self.size):
             self.agents.append(Person(self.house, self, self.ages))
 
     def get_infectious_count(self):
         """Get the number of infectious people in the household."""
-        ic = 0
-        for i in range(0, self.size):
-            if (
-                self.agents[i].status == "infectious"
-                and self.agents[i].hospitalised is False
-            ):
-                ic += 1
-        return ic
+        return len(
+            list(
+                agent
+                for agent in self.agents
+                if agent.status == "infectious" and not agent.hospitalised
+            )
+        )
 
     def is_infected(self):
         """Check if the household has any infectious people."""

facs/base/location.py

@@ -49,7 +49,7 @@ def register_visit(self, e, person, need, deterministic):
         visit_time = self.avg_visit_time
         if person.status == "dead":
             return
-        elif person.status == "infectious":
+        if person.status == "infectious":
             visit_time *= (
                 e.self_isolation_multiplier
             )  # implementing case isolation (CI)

tests/test_household.py

@@ -39,6 +39,30 @@ def test_household_initialization(sample_house, sample_ages):
     assert len(household.agents) == household.size
 
 
+def test_household_initialization_no_size(sample_house, sample_ages):
+    """Test the initialization of Household class."""
+
+    household = Household(sample_house, sample_ages)
+    assert household.house == sample_house
+    assert household.ages == sample_ages
+    assert len(household.agents) == household.size
+    assert household.size in [1, 2, 3, 4]
+
+
+def test_household_initialization_bad_size(sample_house, sample_ages):
+    """Test the initialization of Household class."""
+
+    with pytest.raises(ValueError):
+        Household(sample_house, sample_ages, size=0)
+
+
+def test_household_initialization_large_size(sample_house, sample_ages):
+    """Test the initialization of Household class."""
+
+    with pytest.warns(UserWarning):
+        Household(sample_house, sample_ages, size=5)
+
+
 def test_get_infectious_count_type(sample_house, sample_ages):
     """Test get_infectious_count method in Household class."""