From fb07a1724d3824ea8251920be9c9feef594c4435 Mon Sep 17 00:00:00 2001
From: jurikane <jurikane@gmail.com>
Date: Wed, 1 Oct 2025 13:45:03 +0900
Subject: [PATCH] security update of tornado and increased coverage through
 unittests

---
 .github/workflows/python-app.yml         |   3 +
 pyproject.toml                           |   2 +-
 requirements.txt                         |   2 +-
 src/agents/area.py                       |   3 +
 src/agents/color_cell.py                 |  27 ++--
 src/config/loader.py                     |   5 +-
 src/models/participation_model.py        |  31 +++--
 src/viz/factory.py                       |   2 +-
 tests/factory.py                         |  27 ++--
 tests/test_color_cell.py                 |  70 ++++++++++
 tests/test_conduct_election.py           |   4 +-
 tests/test_create_personalities.py       |   4 +-
 tests/test_distribute_rewards.py         |   4 +-
 tests/test_initialize_all_areas.py       |  14 +-
 tests/test_loader.py                     |  94 +++++++++++++
 tests/test_participation_area_agent.py   | 102 +++++++++++++-
 tests/test_participation_model.py        | 166 ++++++++++++++++-------
 tests/test_participation_voting_agent.py |  13 +-
 tests/test_pers_dist.py                  |  42 +++---
 tests/test_tally_votes.py                |   4 +-
 tests/test_update_color_distribution.py  |  15 +-
 21 files changed, 488 insertions(+), 146 deletions(-)
 create mode 100644 tests/test_color_cell.py
 create mode 100644 tests/test_loader.py

diff --git a/.github/workflows/python-app.yml b/.github/workflows/python-app.yml
index 5edea8b..7819fdf 100644
--- a/.github/workflows/python-app.yml
+++ b/.github/workflows/python-app.yml
@@ -6,6 +6,9 @@ on:
   pull_request:
     branches: [ main ]
 
+permissions:
+  contents: read
+
 jobs:
   build:
 
diff --git a/pyproject.toml b/pyproject.toml
index ad4c30e..a418c0a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -38,7 +38,7 @@ annotated-types = ">=0.6.0"
 typing-inspection = ">=0.4.0"
 PyYAML = "6.0.3"
 toml = "0.10.2"
-tornado = "6.4"
+tornado = ">=6.5.2,<7"
 pytest = "8.2.0"
 pytest-cov = "5.0.0"
 iniconfig = "*"
diff --git a/requirements.txt b/requirements.txt
index f7d1dc6..bf4cd3c 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -17,7 +17,7 @@ annotated-types>=0.6.0
 typing-inspection>=0.4.0
 PyYAML==6.0.3
 toml==0.10.2
-tornado==6.4
+tornado>=6.5.2
 pytest==8.2.0
 pytest-cov==5.0.0
 iniconfig
diff --git a/src/agents/area.py b/src/agents/area.py
index 4ae5b1c..fd10519 100644
--- a/src/agents/area.py
+++ b/src/agents/area.py
@@ -184,6 +184,9 @@ def add_agent(self, agent: VoteAgent) -> None:
         Args:
             agent (VoteAgent): The agent to be added to the area.
         """
+        # Make sure its an instance of Agent
+        if not isinstance(agent, Agent):
+            raise ValueError("Only VoteAgent instances can be added to an Area")
         self.agents.append(agent)
 
     def add_cell(self, cell: ColorCell) -> None:
diff --git a/src/agents/color_cell.py b/src/agents/color_cell.py
index 4032906..080d319 100644
--- a/src/agents/color_cell.py
+++ b/src/agents/color_cell.py
@@ -1,4 +1,4 @@
-from mesa import Agent
+from mesa import Agent, Model
 
 
 class ColorCell(Agent):
@@ -9,7 +9,7 @@ class ColorCell(Agent):
         color (int): The color of the cell.
     """
 
-    def __init__(self, unique_id, model, pos: tuple, initial_color: int):
+    def __init__(self, unique_id: int, model: Model, pos: tuple, initial_color: int):
         """
         Initializes a ColorCell, at the given row, col position.
 
@@ -20,33 +20,30 @@ def __init__(self, unique_id, model, pos: tuple, initial_color: int):
             initial_color (int): The initial color of the cell.
         """
         super().__init__(unique_id, model)
-        # The "pos" variable in mesa is special, so I avoid it here
+        # self.pos will be set by the grid when we place the agent
         self._row = pos[0]
         self._col = pos[1]
         self.color = initial_color  # The cell's current color (int)
         self._next_color = None
-        self.agents = []
-        self.areas = []
+        self.agents = []    # TODO change to using mesas AgentSet class!
+        self.areas = []    # TODO change to using mesas AgentSet class!
         self.is_border_cell = False
+        # Add it to the models grid
+        model.grid.place_agent(self, pos)
 
     def __str__(self):
-        return (f"Cell ({self.unique_id}, pos={self.position}, "
+        return (f"Cell ({self.unique_id}, pos={self.pos}, "
                 f"color={self.color}, num_agents={self.num_agents_in_cell})")
 
-    @property
-    def col(self):
-        """The col location of this cell."""
-        return self._col
-
     @property
     def row(self):
         """The row location of this cell."""
-        return self._row
+        return self.pos[0]
 
     @property
-    def position(self):  # The variable pos is special in mesa!
-        """The location of this cell."""
-        return self._row, self._col
+    def col(self):
+        """The col location of this cell."""
+        return self.pos[1]
 
     @property
     def num_agents_in_cell(self):
diff --git a/src/config/loader.py b/src/config/loader.py
index 9e8256a..07ef733 100644
--- a/src/config/loader.py
+++ b/src/config/loader.py
@@ -13,9 +13,12 @@ def check_schema(open_file):
     return AppConfig.model_validate(raw)
 
 
-def load_config(config_file=None):
+def load_config(config_file=None) -> AppConfig:
     """
     Load configuration from a YAML file.
+
+    Returns:
+        AppConfig: Validated configuration object.
     """
     if config_file is None:
         config_file = os.environ.get("CONFIG_FILE", "default.yaml")
diff --git a/src/models/participation_model.py b/src/models/participation_model.py
index 6c3ae95..62a7eef 100644
--- a/src/models/participation_model.py
+++ b/src/models/participation_model.py
@@ -53,8 +53,8 @@ class ParticipationModel(mesa.Model):
     Attributes:
         grid (mesa.space.SingleGrid): Grid representing the environment
             with a single occupancy per cell (the color).
-        height (int): The height of the grid.
-        width (int): The width of the grid.
+        grid.height (int): The height of the grid.
+        grid.width (int): The width of the grid.
         colors (ndarray): Array containing the unique color identifiers.
         voting_rule (Callable): A function defining the social welfare
             function to aggregate agent preferences. This callable typically
@@ -110,8 +110,6 @@ def __init__(self, height, width, num_agents, num_colors, num_personalities,
         else:
             self.np_random = np.random.default_rng()
         # TODO clean up class (public/private variables)
-        self.height = height
-        self.width = width
         self.colors = np.arange(num_colors)
         # Create a scheduler that goes through areas first then color cells
         self.scheduler = CustomScheduler(self)
@@ -138,21 +136,21 @@ def __init__(self, height, width, num_agents, num_colors, num_personalities,
         # Election impact factor on color mutation through a probability array
         self.color_probs = self.init_color_probs(election_impact_on_mutation)
         # Create search pairs once for faster iterations when comparing rankings
-        self.search_pairs = list(combinations(range(0, self.options.size), 2))  # TODO check if correct!
-        self.option_vec = np.arange(self.options.size)  # Also to speed up
+        self.search_pairs = list(combinations(range(0, self.options.shape[0]), 2))  # TODO check if correct!
+        self.option_vec = np.arange(self.options.shape[0])  # Also to speed up
         self.color_search_pairs = list(combinations(range(0, num_colors), 2))
         # Create color cells (IDs start after areas+agents)
-        self.color_cells: List[Optional[ColorCell]] = [None] * (height * width)
+        self.color_cells: List[Optional[ColorCell]] = [None] * (height * width)  # TODO change to using mesas AgentSet class!
         self._initialize_color_cells(id_start=num_agents + num_areas)
         # Create voting agents (IDs start after areas)
         # TODO: Where do the agents get there known cells from and how!?
-        self.voting_agents: List[Optional[VoteAgent]] = [None] * num_agents
+        self.voting_agents: List[Optional[VoteAgent]] = [None] * num_agents    # TODO change to using mesas AgentSet class!
         self.personalities = self.create_personalities(num_personalities)
         self.personality_distribution = self.pers_dist(num_personalities)
         self.initialize_voting_agents(id_start=num_areas)
         # Area variables
         self.global_area = self.initialize_global_area()  # TODO create bool variable to make this optional
-        self.areas: List[Optional[Area]] = [None] * num_areas
+        self.areas: List[Optional[Area]] = [None] * num_areas    # TODO change to using mesas AgentSet class!
         self.av_area_height = av_area_height
         self.av_area_width = av_area_width
         self.area_size_variance = area_size_variance
@@ -165,6 +163,14 @@ def __init__(self, height, width, num_agents, num_colors, num_personalities,
         # Collect initial data
         self.datacollector.collect(self)
 
+    @property
+    def height(self) -> int:
+        return self.grid.height
+
+    @property
+    def width(self) -> int:
+        return self.grid.width
+
     @property
     def num_colors(self) -> int:
         return len(self.colors)
@@ -203,12 +209,10 @@ def _initialize_color_cells(self, id_start=0) -> None:
             color = self.color_by_dst(self._preset_color_dst)
             # Create the cell (skip ids for area and voting agents)
             cell = ColorCell(unique_id, self, (row, col), color)
-            # Add it to the grid
-            self.grid.place_agent(cell, (row, col))
             # Add the color cell to the scheduler
             #self.scheduler.add(cell) # TODO: check speed diffs using this..
             # And to the 'model.color_cells' list (for faster access)
-            self.color_cells[idx] = cell  # TODO: check if its not better to simply use the grid when finally changing the grid type to SingleGrid
+            self.color_cells[idx] = cell  # TODO: change to using the grid
 
     def initialize_voting_agents(self, id_start=0) -> None:
         """
@@ -219,6 +223,9 @@ def initialize_voting_agents(self, id_start=0) -> None:
         Args:
             id_start (int): The starting ID for agents to ensure unique IDs.
         """
+        # Testing parameter validity
+        if self.num_agents < 1:
+            raise ValueError("The number of agents must be at least 1.")
         dist = self.personality_distribution
         assets = self.common_assets // self.num_agents
         for idx in range(self.num_agents):
diff --git a/src/viz/factory.py b/src/viz/factory.py
index 436e444..9fbc6d2 100644
--- a/src/viz/factory.py
+++ b/src/viz/factory.py
@@ -56,7 +56,7 @@ def portrayal(agent):
             "Layer": 0,
             "Color": color_name,
             # Hover fields:
-            "Position": f"{agent.position}",
+            "Position": f"{agent.pos}",
             "Color - text": color_name,
         }
 
diff --git a/tests/factory.py b/tests/factory.py
index b444326..6d54323 100644
--- a/tests/factory.py
+++ b/tests/factory.py
@@ -1,12 +1,23 @@
 from src.models.participation_model import ParticipationModel
-from pathlib import Path
-import yaml
+from src.model_setup import build_model_kwargs
+from src.config.loader import load_config
+from src.config.schema import AppConfig
 
-DEFAULT_CONFIG = Path("configs/default.yaml")
 
-def create_default_model(**overrides):
-    with open(DEFAULT_CONFIG, "r") as f:
-        config = yaml.safe_load(f)
-    params = config["model"]
+def create_test_model(**overrides) -> tuple[ParticipationModel, dict]:
+    """
+    Create a ParticipationModel instance using the default config
+    (set by DEFAULT_CONFIG, fallback to 'configs/default.yaml'),
+    returning both the model and the parameter dictionary used.
+    This is useful for tests that need to inspect model parameters.
+
+    Args:
+        **overrides: Any model parameters to override from defaults.
+    Returns:
+        tuple: (ParticipationModel instance, config dict)
+    """
+    model_app_cfg = load_config().model
+    params = build_model_kwargs(model_app_cfg)
     params.update(overrides)
-    return ParticipationModel(**params)
+    model = ParticipationModel(**params)
+    return model, params
diff --git a/tests/test_color_cell.py b/tests/test_color_cell.py
new file mode 100644
index 0000000..d76d816
--- /dev/null
+++ b/tests/test_color_cell.py
@@ -0,0 +1,70 @@
+import unittest
+from mesa import Model, space
+from src.agents.color_cell import ColorCell
+
+
+class DummyModel(Model):
+    """Minimal model so ColorCell can be instantiated."""
+    def __init__(self):
+        super().__init__()
+        self.grid = space.SingleGrid(10, 10, torus=True)
+
+
+class TestColorCell(unittest.TestCase):
+
+    def setUp(self):
+        self.model = DummyModel()
+
+    def test_initialization(self):
+        cell = ColorCell(unique_id=1, model=self.model, pos=(2, 3), initial_color=5)
+
+        self.assertEqual(cell.unique_id, 1)
+        self.assertEqual(cell.row, 2)
+        self.assertEqual(cell.col, 3)
+        self.assertEqual(cell.pos, (2, 3))
+        self.assertEqual(cell.color, 5)
+        self.assertEqual(cell.num_agents_in_cell, 0)
+        self.assertEqual(cell.agents, [])
+        self.assertEqual(cell.areas, [])
+        self.assertFalse(cell.is_border_cell)
+
+    def test_add_and_remove_agents(self):
+        cell = ColorCell(1, self.model, (0, 0), 1)
+        agent1, agent2 = object(), object()
+
+        cell.add_agent(agent1)
+        cell.add_agent(agent2)
+
+        self.assertEqual(cell.num_agents_in_cell, 2)
+        self.assertIn(agent1, cell.agents)
+        self.assertIn(agent2, cell.agents)
+
+        cell.remove_agent(agent1)
+        self.assertEqual(cell.num_agents_in_cell, 1)
+        self.assertNotIn(agent1, cell.agents)
+
+    def test_add_area(self):
+        cell = ColorCell(1, self.model, (0, 0), 1)
+        area = {"id": 123}
+
+        cell.add_area(area)
+        self.assertIn(area, cell.areas)
+
+    def test_str_representation(self):
+        cell = ColorCell(42, self.model, (1, 1), 7)
+        s = str(cell)
+        self.assertIn("Cell (42", s)
+        self.assertIn("pos=(1, 1)", s)
+        self.assertIn("color=7", s)
+        self.assertIn("num_agents=0", s)
+
+    # TODO: Add tests for color_step once implemented
+    #  - Check that it picks a next color from neighbors
+    #  - Handle ties correctly
+    #
+    # TODO: Add tests for advance once implemented
+    #  - Ensure that color updates from _next_color
+
+
+if __name__ == "__main__":
+    unittest.main()
\ No newline at end of file
diff --git a/tests/test_conduct_election.py b/tests/test_conduct_election.py
index 38a883a..f77a25d 100644
--- a/tests/test_conduct_election.py
+++ b/tests/test_conduct_election.py
@@ -1,12 +1,12 @@
 import unittest
 from unittest.mock import MagicMock
-from tests.factory import create_default_model
+from tests.factory import create_test_model
 
 # TODO add more complex tests
 
 class TestConductElection(unittest.TestCase):
     def setUp(self):
-        self.model = create_default_model(num_areas=1)
+        self.model, _ = create_test_model(num_areas=1)
         self.model.initialize_area = MagicMock()
 
     def test_election_returns_integer_turnout(self):
diff --git a/tests/test_create_personalities.py b/tests/test_create_personalities.py
index 3bfa193..daacef4 100644
--- a/tests/test_create_personalities.py
+++ b/tests/test_create_personalities.py
@@ -1,7 +1,7 @@
 import unittest
 from math import factorial
 from itertools import permutations
-from tests.factory import create_default_model
+from tests.factory import create_test_model
 from unittest.mock import MagicMock
 
 
@@ -9,7 +9,7 @@ class TestParticipationModel(unittest.TestCase):
 
     def setUp(self):
         """Create a fresh model instance before each test and mock `initialize_area`."""
-        self.model = create_default_model(
+        self.model, _ = create_test_model(
             height=10, width=10, num_agents=100, num_colors=4,
             num_personalities=10, area_size_variance=0.2,
             num_areas=4, av_area_height=5, av_area_width=5,
diff --git a/tests/test_distribute_rewards.py b/tests/test_distribute_rewards.py
index 4197a3e..e0f771d 100644
--- a/tests/test_distribute_rewards.py
+++ b/tests/test_distribute_rewards.py
@@ -1,10 +1,10 @@
 import unittest
 from unittest.mock import MagicMock
-from tests.factory import create_default_model
+from tests.factory import create_test_model
 
 class TestDistributeRewards(unittest.TestCase):
     def setUp(self):
-        self.model = create_default_model(num_areas=1)
+        self.model, _ = create_test_model(num_areas=1)
         self.model.initialize_area = MagicMock()
 
     def test_distribute(self):
diff --git a/tests/test_initialize_all_areas.py b/tests/test_initialize_all_areas.py
index bea2faa..d46e568 100644
--- a/tests/test_initialize_all_areas.py
+++ b/tests/test_initialize_all_areas.py
@@ -1,14 +1,14 @@
 import unittest
 from unittest.mock import MagicMock
 from numpy import sqrt
-from tests.factory import create_default_model  # Import from factory.py
+from tests.factory import create_test_model  # Import from factory.py
 
 
 class TestParticipationModelInitializeAllAreas(unittest.TestCase):
 
     def setUp(self):
         """Create a fresh model instance before each test and mock `initialize_area`."""
-        self.model = create_default_model(
+        self.model, _ = create_test_model(
             num_areas=4,  # Override num_areas to 4
             height=10,  # Set grid height
             width=10,  # Set grid width
@@ -31,7 +31,7 @@ def test_initialize_all_areas_uniform_distribution(self):
 
     def test_initialize_all_areas_with_non_square_number(self):
         """Test that the method handles non-square numbers by adding extra areas randomly."""
-        model = create_default_model(
+        model, _ = create_test_model(
             num_areas=5,  # Override num_areas to 5
         )
         # model.initialize_all_areas()  # Runs on initialization
@@ -40,14 +40,14 @@ def test_initialize_all_areas_with_non_square_number(self):
 
     def test_initialize_all_areas_no_areas(self):
         """Test that the method does nothing if num_areas is 0."""
-        model = create_default_model(
+        model, _ = create_test_model(
             num_areas=0,  # Set num_areas to 0
         )
         assert model.num_areas == 0  # Verify no areas were initialized
 
     def test_initialize_all_areas_random_additional_areas(self):
         """Test that additional areas are placed randomly if num_areas exceeds uniform grid capacity."""
-        model = create_default_model(
+        model, _ = create_test_model(
             num_areas=5,  # Override num_areas to 5
             height=10,
             width=10,
@@ -68,7 +68,7 @@ def test_initialize_all_areas_random_additional_areas(self):
 
     def test_initialize_all_areas_handles_non_square_distribution(self):
         """Test that the number of areas matches `num_areas` even for non-square cases."""
-        model = create_default_model(
+        model, _ = create_test_model(
             num_areas=6,  # Override num_areas to 6
         )
         # Check that exactly 6 areas are initialized
@@ -76,7 +76,7 @@ def test_initialize_all_areas_handles_non_square_distribution(self):
 
     def test_initialize_all_areas_calculates_distances_correctly(self):
         """Test that area distances are calculated correctly."""
-        model = create_default_model(
+        model, _ = create_test_model(
             num_areas=4,  # Override num_areas to 4
             height=10,
             width=10,
diff --git a/tests/test_loader.py b/tests/test_loader.py
new file mode 100644
index 0000000..ae4c7c8
--- /dev/null
+++ b/tests/test_loader.py
@@ -0,0 +1,94 @@
+import pytest
+import tempfile
+import os
+from pathlib import Path
+from src.config import loader
+from unittest.mock import patch
+
+@pytest.fixture
+def mock_model_validate():
+    with patch("src.config.schema.AppConfig.model_validate",
+               return_value="validated") as mock:
+        yield mock
+
+def test_check_schema_valid(tmp_path):
+    cfg = tmp_path / "cfg.yaml"
+    cfg.write_text("key: value")
+    with cfg.open("r") as f:
+        with patch("src.config.schema.AppConfig.model_validate", return_value="ok") as mock:
+            result = loader.check_schema(f)
+    assert result == "ok"
+    mock.assert_called_once()
+
+def test_load_config_default_yaml(monkeypatch, tmp_path):
+    default = tmp_path / "default.yaml"
+    default.write_text("key: value")
+    monkeypatch.chdir(tmp_path)
+    with patch("src.config.schema.AppConfig.model_validate", return_value="validated"):
+        result = loader.load_config()
+    assert result == "validated"
+
+def test_load_config_direct_path(mock_model_validate):
+    with tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False) as tmp:
+        tmp.write("key: value")
+        tmp_path = Path(tmp.name)
+    result = loader.load_config(str(tmp_path))
+    assert result == "validated"
+    tmp_path.unlink()
+
+def test_load_config_cwd(mock_model_validate):
+    with tempfile.NamedTemporaryFile("w", suffix=".yaml", dir=Path.cwd(), delete=False) as tmp:
+        tmp.write("key: value")
+        tmp_path = Path(tmp.name)
+    result = loader.load_config(tmp_path.name)
+    assert result == "validated"
+    tmp_path.unlink()
+
+def test_load_config_project_root(mock_model_validate, tmp_path, monkeypatch):
+    # Simulate project root two levels up
+    project_root = tmp_path / "project"
+    configs_dir = project_root / "configs"
+    configs_dir.mkdir(parents=True)
+    cfg_file = configs_dir / "test.yaml"
+    cfg_file.write_text("key: value")
+
+    # Patch __file__ to trick loader into thinking it's inside project_root/src/config
+    fake_loader_file = project_root / "src" / "config" / "loader.py"
+    fake_loader_file.parent.mkdir(parents=True)
+    fake_loader_file.write_text("# fake loader file")
+
+    monkeypatch.setattr(loader, "__file__", str(fake_loader_file))
+
+    result = loader.load_config("test.yaml")
+    assert result == "validated"
+
+def test_load_config_legacy(mock_model_validate, tmp_path, monkeypatch):
+    # Create fake project root
+    project_root = tmp_path / "project"
+    legacy_dir = project_root / "configs"
+    legacy_dir.mkdir(parents=True)
+    cfg_file = legacy_dir / "legacy.yaml"
+    cfg_file.write_text("key: value")
+
+    # Fake loader's __file__ so parents[2] == project_root
+    fake_loader_file = project_root / "tmp" / "config" / "loader.py"
+    fake_loader_file.parent.mkdir(parents=True)
+    fake_loader_file.write_text("# fake loader file")
+    monkeypatch.setattr(loader, "__file__", str(fake_loader_file))
+    # Load
+    result = loader.load_config("legacy.yaml")
+    assert result == "validated"
+
+def test_load_config_env_var(mock_model_validate):
+    with tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False) as tmp:
+        tmp.write("key: value")
+        tmp_path = Path(tmp.name)
+    os.environ["CONFIG_FILE"] = str(tmp_path)
+    result = loader.load_config()
+    assert result == "validated"
+    tmp_path.unlink()
+    del os.environ["CONFIG_FILE"]
+
+def test_load_config_not_found():
+    with pytest.raises(FileNotFoundError):
+        loader.load_config("nonexistent.yaml")
diff --git a/tests/test_participation_area_agent.py b/tests/test_participation_area_agent.py
index 3b1fcb9..2131fb2 100644
--- a/tests/test_participation_area_agent.py
+++ b/tests/test_participation_area_agent.py
@@ -1,19 +1,107 @@
 import unittest
 import random
 import numpy as np
-from src.models.participation_model import Area
+from src.agents.area import Area
+from src.agents.color_cell import ColorCell
 from src.agents.vote_agent import VoteAgent
-from .test_participation_model import TestParticipationModel, model_cfg
+from src.model_setup import build_model_kwargs
+from src.models.participation_model import ParticipationModel
+from src.config.loader import load_config
 from src.utils.social_welfare_functions import majority_rule, approval_voting
 from src.utils.distance_functions import kendall_tau, spearman
+import mesa.space as space
+from mesa import Model
 
 
-class TestArea(unittest.TestCase):
+###################################
+# Dummy classes for isolated tests #
+###################################
+
+class DummyModel(Model):
+    """Minimal model for testing Area without full ParticipationModel."""
+    def __init__(self, width=5, height=5, num_colors=3):
+        super().__init__()
+        self.width = width
+        self.height = height
+        self.num_colors = num_colors
+        self.grid = space.SingleGrid(height=height, width=width, torus=True)
+        self.personalities = [0, 1]
+        self.voting_agents = []
+        # Stubs for election-related attributes
+        self.distance_func = lambda *args, **kwargs: 0
+        self.voting_rule = lambda x: [0]
+        self.options = [0, 1, 2]
+        self.known_cells = 1
+        self.color_search_pairs = []
+        self.max_reward = 10
+        self.election_costs = 1
+        self.mu = 0.1
+        self.color_probs = [1 / num_colors] * num_colors
+        self.np_random = np.random.default_rng()
+
+##################################
+# Unit tests for the Area class  #
+##################################
+
+class TestAreaBasics(unittest.TestCase):
+
+    def setUp(self):
+        self.dummy_model = DummyModel()
+
+    def test_initialization_no_variance(self):
+        area = Area(unique_id=1, model=self.dummy_model, height=2, width=3, size_variance=0)
+        self.assertEqual(area.num_cells, 6)
+        self.assertEqual(area.num_agents, 0)
+        self.assertTrue((area.color_distribution == np.zeros(self.dummy_model.num_colors)).all())
+
+    def test_invalid_variance_raises(self):
+        with self.assertRaises(ValueError):
+            Area(unique_id=2, model=self.dummy_model, height=2, width=3, size_variance=1.5)
+
+    def test_add_agent_and_cell(self):
+        area = Area(1, self.dummy_model, 2, 2, 0)
+        cell = ColorCell(10, self.dummy_model, (0, 0), 1)
+        area.add_cell(cell)
+        dummy_agent = VoteAgent(1, self.dummy_model, (0, 0))
+        area.add_agent(dummy_agent)
+
+        self.assertIn(cell, area.cells)
+        self.assertIn(dummy_agent, area.agents)
+        self.assertEqual(area.num_agents, 1)
+
+    def test_idx_field_assigns_cells(self):
+        # Replace the automatically placed cell with our own
+        cell = ColorCell(11, self.dummy_model, (0, 0), 2)
+        present_cell = self.dummy_model.grid.get_cell_list_contents([(0, 0)])[0]
+        self.dummy_model.grid.remove_agent(present_cell)
+        self.dummy_model.grid.place_agent(cell, (0, 0))
+
+        area = Area(1, self.dummy_model, 1, 1, 0)
+        area.idx_field = (0, 0)
+
+        self.assertEqual(area.idx_field, (0, 0))
+        self.assertIn(cell, area.cells)
+        self.assertAlmostEqual(area.color_distribution.sum(), 1.0, places=7)
+
+    def test_str_representation(self):
+        area = Area(99, self.dummy_model, 2, 3, 0)
+        s = str(area)
+        self.assertIn("Area(id=99", s)
+        self.assertIn("size=2x3", s)
+        self.assertIn("num_agents=0", s)
+        self.assertIn("num_cells=6", s)
+
+
+################################################################
+# Integration tests for Area within ParticipationModel context #
+################################################################
+
+class TestAreaIntegration(unittest.TestCase):
 
     def setUp(self):
-        test_model = TestParticipationModel()
-        test_model.setUp()
-        self.model = test_model.model
+        self.model_cfg = load_config().model
+        model_cfg = build_model_kwargs(self.model_cfg)
+        self.model = ParticipationModel(**model_cfg)
 
     def test_update_color_distribution(self):
         rand_area = random.sample(self.model.areas, 1)[0]
@@ -68,7 +156,7 @@ def test_conduct_election(self):
     def test_adding_new_area_and_agent_within_it(self):
         # Additional area and agent
         personality = random.choice(self.model.personalities)
-        a = VoteAgent(model_cfg["num_agents"] + 1, self.model, pos=(0, 0),
+        a = VoteAgent(self.model_cfg.num_agents + 1, self.model, pos=(0, 0),
                       personality=personality, assets=25)
         additional_test_area = Area(self.model.num_areas + 1,
                                     model=self.model, height=5,
diff --git a/tests/test_participation_model.py b/tests/test_participation_model.py
index 6a697dc..d69bf14 100644
--- a/tests/test_participation_model.py
+++ b/tests/test_participation_model.py
@@ -1,88 +1,158 @@
-import unittest
-from src.models.participation_model import (ParticipationModel, Area,
-                                            distance_functions,
-                                            social_welfare_functions)
-from src.config.loader import load_config
 import mesa
+import unittest
+import numpy as np
+from src.models.participation_model import (
+    ParticipationModel, Area,
+    distance_functions,
+    social_welfare_functions
+)
+from tests.factory import create_test_model
 
-config = load_config()
-model_cfg = config.model.model_dump()
-vis_cfg = config.visualization.model_dump()
 
+class TestParticipationModelUnit(unittest.TestCase):
+    def setUp(self):
+        self.model, self.model_cfg = create_test_model()
 
-class TestParticipationModel(unittest.TestCase):
+    ###################################
+    # Basic unit tests for Area class #
+    ###################################
 
-    def setUp(self):
-        self.model = ParticipationModel(**model_cfg)
+    def test_initialization_creates_expected_components(self):
+        model_cfg = self.model_cfg
+        self.assertEqual(self.model.grid.width, model_cfg["width"])
+        self.assertEqual(self.model.grid.height, model_cfg["height"])
+        self.assertEqual(len(self.model.voting_agents), model_cfg["num_agents"])
+        self.assertEqual(len(self.model.color_cells),
+                         model_cfg["width"] * model_cfg["height"])
+        self.assertEqual(len(self.model.areas), model_cfg["num_areas"])
+        self.assertIsNotNone(self.model.global_area)
+
+    def test_personality_distribution_sums_to_one(self):
+        dist = self.model.personality_distribution
+        np.testing.assert_almost_equal(dist.sum(), 1.0)
+        self.assertEqual(len(dist), len(self.model.personalities))
+
+    def test_preset_color_distribution_valid(self):
+        dst = self.model.preset_color_dst
+        np.testing.assert_almost_equal(sum(dst), 1.0)
+        self.assertTrue(all(p >= 0 for p in dst))
+
+    # --- Static & helper methods ---
+
+    def test_color_by_dst_respects_distribution(self):
+        probs = np.array([0.1, 0.3, 0.6])
+        counts = [0, 0, 0]
+        for _ in range(1000):
+            c = ParticipationModel.color_by_dst(probs)
+            counts[c] += 1
+        self.assertGreater(counts[2], counts[1])
+        self.assertGreater(counts[1], counts[0])
+
+    def test_create_all_options_without_ties(self):
+        opts = ParticipationModel.create_all_options(3)
+        self.assertEqual(opts.shape[1], 3)
+        for row in opts:
+            self.assertEqual(sorted(row), [0, 1, 2])
+
+    def test_create_all_options_with_ties(self):
+        opts = ParticipationModel.create_all_options(2, include_ties=True)
+        self.assertIsInstance(opts, np.ndarray)
+        self.assertEqual(opts.shape[1], 2)
+
+    def test_pers_dist_sums_to_one(self):
+        dist = ParticipationModel.pers_dist(5)
+        np.testing.assert_almost_equal(dist.sum(), 1.0)
+
+    # --- Functional behavior ---
+
+    def test_step_updates_model(self):
+        before = self.model.av_area_color_dst.copy()
+        self.model.step()
+        after = self.model.av_area_color_dst
+        self.assertEqual(len(before), len(after))
+        np.testing.assert_almost_equal(after.sum(), 1.0, decimal=6)
+
+    def test_update_av_area_color_dst(self):
+        self.model.update_av_area_color_dst()
+        dst = self.model.av_area_color_dst
+        np.testing.assert_almost_equal(dst.sum(), 1.0)
 
-    # def test_empty_model(self):
-    #     # TODO: Test empty model
-    #     model = ParticipationModel(10, 10, 0, 1, 0, 1, 0, 1, 1, 0.1, 1, 0, False, 1, 1, 1, 1, 1, False)
-    #     self.assertEqual(model.num_agents, 0)
+    def test_init_color_probs(self):
+        probs = self.model.init_color_probs(1.0)
+        self.assertEqual(probs.shape, (self.model.num_colors,))
+        np.testing.assert_almost_equal(probs.sum(), 1.0)
+
+    def test_initialize_area_adds_area(self):
+        old_num = sum(a is not None for a in self.model.areas)
+        self.model.initialize_area(0, 0, 0)
+        new_num = sum(a is not None for a in self.model.areas)
+        self.assertGreaterEqual(new_num, old_num)
+
+    ################################################################
+    # Integration tests for Area within ParticipationModel context #
+    ################################################################
 
     def test_initialization(self):
-        areas_count = len([area for area in self.model.areas
-                           if isinstance(area, Area)])
+        areas_count = len([
+            area for area in self.model.areas if isinstance(area, Area)])
         self.assertEqual(areas_count, self.model.num_areas)
         self.assertIsInstance(self.model.datacollector, mesa.DataCollector)
-        # TODO ... more tests
 
     def test_model_options(self):
-        self.assertEqual(self.model.num_agents, model_cfg["num_agents"])
-        self.assertEqual(self.model.num_colors, model_cfg["num_colors"])
-        self.assertEqual(self.model.num_areas, model_cfg["num_areas"])
+        self.assertEqual(self.model.num_agents, self.model_cfg["num_agents"])
+        self.assertEqual(self.model.num_colors, self.model_cfg["num_colors"])
+        self.assertEqual(self.model.num_areas, self.model_cfg["num_areas"])
         self.assertEqual(self.model.area_size_variance,
-                         model_cfg["area_size_variance"])
-        v_rule = social_welfare_functions[model_cfg["rule_idx"]]
-        dist_func = distance_functions[model_cfg["distance_idx"]]
-        self.assertEqual(self.model.common_assets, model_cfg["common_assets"])
+                         self.model_cfg["area_size_variance"])
+
+        v_rule = social_welfare_functions[self.model_cfg["rule_idx"]]
+        dist_func = distance_functions[self.model_cfg["distance_idx"]]
+
+        self.assertEqual(self.model.common_assets,
+                         self.model_cfg["common_assets"])
         self.assertEqual(self.model.voting_rule, v_rule)
         self.assertEqual(self.model.distance_func, dist_func)
-        self.assertEqual(self.model.election_costs, model_cfg["election_costs"])
+        self.assertEqual(self.model.election_costs,
+                         self.model_cfg["election_costs"])
 
     def test_create_color_distribution(self):
         eq_dst = self.model.create_color_distribution(heterogeneity=0)
-        self.assertEqual([1/model_cfg["num_colors"] for _ in eq_dst], eq_dst)
-        print(f"Color distribution with heterogeneity=0: {eq_dst}")
+        np.testing.assert_allclose(
+            eq_dst, [1 / self.model_cfg["num_colors"]] * len(eq_dst))
+
         het_dst = self.model.create_color_distribution(heterogeneity=1)
-        print(f"Color distribution with heterogeneity=1: {het_dst}")
         mid_dst = self.model.create_color_distribution(heterogeneity=0.5)
-        print(f"Color distribution with heterogeneity=0.5: {mid_dst}")
-        assert het_dst != eq_dst
-        assert mid_dst != eq_dst
-        assert het_dst != mid_dst
+
+        self.assertFalse(np.allclose(het_dst, eq_dst))
+        self.assertFalse(np.allclose(mid_dst, eq_dst))
+        self.assertFalse(np.allclose(het_dst, mid_dst))
 
     def test_distribution_of_personalities(self):
         p_dist = self.model.personality_distribution
-        self.assertAlmostEqual(sum(p_dist), 1.0)
-        self.assertEqual(len(p_dist), model_cfg["num_personalities"])
+        self.assertAlmostEqual(float(sum(p_dist)), 1.0)
+        self.assertEqual(len(p_dist), self.model_cfg["num_personalities"])
+
         voting_agents = self.model.voting_agents
         nr_agents = self.model.num_agents
         personalities = list(self.model.personalities)
         p_counts = {str(i): 0 for i in personalities}
-        # Count the occurrence of each personality
+
         for agent in voting_agents:
             p_counts[str(agent.personality)] += 1
-        # Normalize the counts to get the real personality distribution
+
         real_dist = [p_counts[str(p)] / nr_agents for p in personalities]
-        # Simple tests
+
         self.assertEqual(len(real_dist), len(p_dist))
         self.assertAlmostEqual(float(sum(real_dist)), 1.0)
-        # Compare each value
-        my_delta = 0.4 / model_cfg["num_personalities"]  # The more personalities, the smaller the delta
+
+        my_delta = 0.4 / self.model_cfg["num_personalities"]
         for p_dist_val, real_p_dist_val in zip(p_dist, real_dist):
             self.assertAlmostEqual(p_dist_val, real_p_dist_val, delta=my_delta)
 
-
     def test_initialize_areas(self):
         # TODO (very non-trivial) - has been tested manually so far.
         pass
 
     def test_step(self):
-        pass
-    # TODO add test_step
-    # def test_step(self):
-    #     initial_data = self.model.datacollector.get_model_vars_dataframe().copy()
-    #     self.model.step()
-    #     new_data = self.model.datacollector.get_model_vars_dataframe()
-    #     self.assertNotEqual(initial_data, new_data)
+        # TODO: Add full step integration test
+        pass
\ No newline at end of file
diff --git a/tests/test_participation_voting_agent.py b/tests/test_participation_voting_agent.py
index 8a3888b..608a7fb 100644
--- a/tests/test_participation_voting_agent.py
+++ b/tests/test_participation_voting_agent.py
@@ -1,24 +1,17 @@
-from .test_participation_model import TestParticipationModel
 import unittest
 from src.models.participation_model import Area
 from src.agents.vote_agent import VoteAgent, combine_and_normalize
-from src.config.loader import load_config
+from tests.factory import create_test_model
 import numpy as np
 import random
 
 
-config = load_config()
-model_cfg = config.model
-vis_cfg = config.visualization
-
 class TestVotingAgent(unittest.TestCase):
 
     def setUp(self):
-        test_model = TestParticipationModel()
-        test_model.setUp()
-        self.model = test_model.model
+        self.model, self.model_cfg = create_test_model()
         personality = random.choice(self.model.personalities)
-        self.agent = VoteAgent(model_cfg.num_agents + 1, self.model,
+        self.agent = VoteAgent(self.model_cfg["num_agents"] + 1, self.model,
                                pos=(0, 0), personality=personality, assets=25)
         self.additional_test_area = Area(self.model.num_areas + 1,
                                          model=self.model, height=5,
diff --git a/tests/test_pers_dist.py b/tests/test_pers_dist.py
index 7b2950d..d85d4a9 100644
--- a/tests/test_pers_dist.py
+++ b/tests/test_pers_dist.py
@@ -16,28 +16,28 @@ def create_gaussian_distribution(size: int) -> np.ndarray:
     #     dist[idx] += 1 - sm
     return dist
 
-# Example usage
-if __name__ == "__main__":
-    nr_options = 20
-    gaussian_dist = create_gaussian_distribution(nr_options)
-    s = gaussian_dist.sum()
 
-    nr_zeroes = gaussian_dist.size - np.count_nonzero(gaussian_dist)
-    print("There are", nr_zeroes, "zero values in the distribution")
-
-    # Plot the distribution
-    plt.plot(gaussian_dist)
-    plt.title("Normalized Gaussian Distribution")
-    plt.show()
-
-    sample_size = 800
-    pool = np.arange(nr_options)
-    rng = np.random.default_rng()
-    print(pool.shape)
-    chosen = rng.choice(pool, sample_size, p=gaussian_dist)
-
-    plt.hist(chosen)
-    plt.show()
+# # Example usage and manual testing
+# nr_options = 20
+# gaussian_dist = create_gaussian_distribution(nr_options)
+# s = gaussian_dist.sum()
+#
+# nr_zeroes = gaussian_dist.size - np.count_nonzero(gaussian_dist)
+# print("There are", nr_zeroes, "zero values in the distribution")
+#
+# # Plot the distribution
+# plt.plot(gaussian_dist)
+# plt.title("Normalized Gaussian Distribution")
+# plt.show()
+#
+# sample_size = 800
+# pool = np.arange(nr_options)
+# rng = np.random.default_rng()
+# print(pool.shape)
+# chosen = rng.choice(pool, sample_size, p=gaussian_dist)
+#
+# plt.hist(chosen)
+# plt.show()
 
 
 def test_distribution_normalized():
diff --git a/tests/test_tally_votes.py b/tests/test_tally_votes.py
index e7f4eda..20bd7c5 100644
--- a/tests/test_tally_votes.py
+++ b/tests/test_tally_votes.py
@@ -1,11 +1,11 @@
 import unittest
 import numpy as np
 from unittest.mock import MagicMock
-from tests.factory import create_default_model
+from tests.factory import create_test_model
 
 class TestTallyVotes(unittest.TestCase):
     def setUp(self):
-        self.model = create_default_model(num_areas=1)
+        self.model, _ = create_test_model(num_areas=1)
         self.model.initialize_area = MagicMock()
 
     def test_tally_votes_array(self):
diff --git a/tests/test_update_color_distribution.py b/tests/test_update_color_distribution.py
index f144409..67dac3e 100644
--- a/tests/test_update_color_distribution.py
+++ b/tests/test_update_color_distribution.py
@@ -1,11 +1,11 @@
 import unittest
 import numpy as np
 from unittest.mock import MagicMock
-from tests.factory import create_default_model
+from tests.factory import create_test_model
 
 class TestUpdateColorDistribution(unittest.TestCase):
     def setUp(self):
-        self.model = create_default_model(
+        self.model, _ = create_test_model(
             num_areas=1,
             num_colors=3
         )
@@ -14,11 +14,14 @@ def setUp(self):
     def test_color_distribution(self):
         area = self.model.areas[0]
         old_dist = np.copy(area._color_distribution)
-        # Manually change some cell colors
-        for cell in area.cells[:3]:
+        # Force all cells to color 1
+        for cell in area.cells:
             cell.color = 1
         area._update_color_distribution()
         new_dist = area._color_distribution
-        # TODO: This test fails sometimes (11.09.25)
-        self.assertFalse(np.array_equal(old_dist, new_dist))  # <--- here
+        # Assert that distribution has changed
+        self.assertFalse(np.array_equal(old_dist, new_dist))
+        # Assert it's a proper probability distribution
         self.assertAlmostEqual(np.sum(new_dist), 1.0, places=5)
+        # Stronger: check all mass on color 1
+        self.assertTrue(np.isclose(new_dist[1], 1.0))