Refactor Agency Configuration File Handling and Add Tests

nalgonda/models/agency_config.py

@@ -1,12 +1,8 @@
-import json
-from pathlib import Path
-
 from agency_swarm import Agent
 from pydantic import BaseModel, Field
 
-from nalgonda.agency_config_lock_manager import AgencyConfigLockManager
-from nalgonda.constants import CONFIG_FILE_BASE, DEFAULT_CONFIG_FILE
 from nalgonda.models.agent_config import AgentConfig
+from nalgonda.persistence.agency_config_file_storage import AgencyConfigFileStorage
 
 
 class AgencyConfig(BaseModel):
@@ -26,20 +22,14 @@ def update_agent_ids_in_config(self, agents: list[Agent]) -> None:
 
     @classmethod
     def load(cls, agency_id: str) -> "AgencyConfig":
-        """Load agency config from file"""
-        config_file_path = cls.get_config_path(agency_id)
-        if not config_file_path.is_file():
-            config_file_path = DEFAULT_CONFIG_FILE
+        """Load agency config from the storage"""
+        with AgencyConfigFileStorage(agency_id) as config_file:
+            config = config_file.load()
 
-        with AgencyConfigLockManager.get_lock(agency_id), config_file_path.open() as file:
-            config = json.load(file)
-            config["agency_id"] = agency_id
-            return cls.model_validate(config)
+        config["agency_id"] = agency_id
+        return cls.model_validate(config)
 
     def save(self) -> None:
-        with AgencyConfigLockManager.get_lock(self.agency_id), self.get_config_path(self.agency_id).open("w") as file:
-            file.write(self.model_dump_json(indent=2))
-
-    @staticmethod
-    def get_config_path(agency_id: str) -> Path:
-        return CONFIG_FILE_BASE.with_name(f"config_{agency_id}.json")
+        """Save agency config to the storage"""
+        with AgencyConfigFileStorage(self.agency_id) as config_file:
+            config_file.save(self.model_dump())

nalgonda/persistence/agency_config_file_storage.py

@@ -0,0 +1,84 @@
+import json
+from pathlib import Path
+
+from nalgonda.constants import CONFIG_FILE_BASE, DEFAULT_CONFIG_FILE
+from nalgonda.persistence.agency_config_lock_manager import AgencyConfigLockManager
+from nalgonda.persistence.agency_config_storage_interface import AgencyConfigStorageInterface
+
+
+class AgencyConfigFileStorage(AgencyConfigStorageInterface):
+    """
+    A thread-safe context manager for handling agency-specific configuration files.
+
+    This class ensures that file operations on configuration files are managed
+    in a thread-safe manner using locks. Each agency identified by its unique ID
+    gets its own lock to prevent concurrent access issues.
+    """
+
+    def __init__(self, agency_id: str):
+        self.config_file_path = self._get_config_path(agency_id)
+        self.lock = None
+        self.agency_id = agency_id
+
+    def __enter__(self):
+        """
+        Enters the runtime context and acquires the lock for the agency file.
+
+        If the specific agency configuration file does not exist, it falls back
+        to a default configuration file.
+
+        Returns:
+            self: An instance of AgencyConfigFileStorage.
+        """
+        self.lock = AgencyConfigLockManager.get_lock(self.agency_id)
+        self.lock.acquire()
+        if not self.config_file_path.is_file():
+            self._create_default_config()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """
+        Exits the runtime context and releases the lock for the agency file.
+        """
+        self.lock.release()
+
+    def load(self):
+        """
+        Loads the configuration from the agency file.
+
+        Returns:
+            dict: The loaded configuration data.
+        """
+        with self.config_file_path.open() as file:
+            return json.load(file)
+
+    def save(self, data):
+        """
+        Saves the provided data to the agency file.
+
+        Args:
+            data (dict): The configuration data to be saved.
+        """
+        with self.config_file_path.open("w") as file:
+            json.dump(data, file, indent=2)
+
+    def _create_default_config(self):
+        """
+        Creates a new configuration file based on the default configuration.
+        """
+        with DEFAULT_CONFIG_FILE.open() as file:
+            config = json.load(file)
+        self.save(config)
+
+    @staticmethod
+    def _get_config_path(agency_id: str) -> Path:
+        """
+        Generates the path for the agency configuration file.
+
+        Args:
+            agency_id (str): The unique identifier for the agency.
+
+        Returns:
+            Path: The path object for the agency configuration file.
+        """
+        return CONFIG_FILE_BASE.with_name(f"config_{agency_id}.json")

nalgonda/persistence/agency_config_storage_interface.py

@@ -0,0 +1,25 @@
+from abc import ABC, abstractmethod
+
+
+class AgencyConfigStorageInterface(ABC):
+    """Interface for agency config storage"""
+
+    @abstractmethod
+    def __enter__(self):
+        """Enter context manager"""
+        pass
+
+    @abstractmethod
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Exit context manager"""
+        pass
+
+    @abstractmethod
+    def load(self):
+        """Load agency config from the storage"""
+        pass
+
+    @abstractmethod
+    def save(self, data):
+        """Save agency config to the storage"""
+        pass

pyproject.toml

@@ -4,6 +4,7 @@ version = "1.0.0"
 description = "Nalgonda is a FastAPI app to manage swarm agencies"
 authors = [
     "Nikita Bobrovskiy <[email protected]>",
+    "Guilherme Parpinelli <[email protected]>",
 ]
 readme = "README.md"
 

tests/persistence/test_agency_config_file_storage.py

@@ -0,0 +1,72 @@
+import json
+from unittest.mock import patch
+
+import pytest
+
+from nalgonda.constants import DEFAULT_CONFIG_FILE
+from nalgonda.persistence.agency_config_file_storage import AgencyConfigFileStorage
+from nalgonda.persistence.agency_config_lock_manager import AgencyConfigLockManager
+
+
+@pytest.fixture
+def mock_config_path(temp_dir):
+    """Fixture to patch the _get_config_path method in AgencyConfigFileStorage."""
+    agency_id = "test_agency"
+    config_path = temp_dir / f"config_{agency_id}.json"
+
+    with patch.object(AgencyConfigFileStorage, "_get_config_path", return_value=config_path):
+        yield config_path, agency_id
+
+
+def test_lock_acquisition_and_release(mock_config_path):
+    _, agency_id = mock_config_path
+    storage = AgencyConfigFileStorage(agency_id)
+    lock = AgencyConfigLockManager.get_lock(agency_id)
+
+    with storage:
+        assert lock.locked(), "Lock was not acquired"
+
+    assert not lock.locked(), "Lock was not released"
+
+
+def test_load_configuration(mock_config_path):
+    config_path, agency_id = mock_config_path
+    config_data = {"key": "value"}
+    config_path.write_text(json.dumps(config_data))
+
+    storage = AgencyConfigFileStorage(agency_id)
+    with storage:
+        loaded_config = storage.load()
+
+    assert loaded_config == config_data, "Loaded configuration does not match expected data"
+
+
+def test_save_configuration(mock_config_path):
+    config_path, agency_id = mock_config_path
+    new_config = {"new_key": "new_value"}
+
+    storage = AgencyConfigFileStorage(agency_id)
+    with storage:
+        storage.save(new_config)
+
+    assert config_path.exists(), "Configuration file was not created"
+    with open(config_path) as file:
+        saved_data = json.load(file)
+    assert saved_data == new_config, "Saved data does not match"
+
+
+def test_default_configuration_used(mock_config_path):
+    config_path, agency_id = mock_config_path
+
+    with open(DEFAULT_CONFIG_FILE) as file:
+        default_config = json.load(file)
+
+    # Ensure the specific config file does not exist to trigger default config usage
+    if config_path.exists():
+        config_path.unlink()
+
+    storage = AgencyConfigFileStorage(agency_id)
+    with storage:
+        loaded_config = storage.load()
+
+    assert loaded_config == default_config, "Default configuration was not used"

tests/persistence/test_agency_config_lock_manager.py

@@ -0,0 +1,43 @@
+import threading
+
+from nalgonda.persistence.agency_config_lock_manager import AgencyConfigLockManager
+
+
+def test_lock_uniqueness():
+    agency_id1 = "agency1"
+    agency_id2 = "agency2"
+
+    lock1 = AgencyConfigLockManager.get_lock(agency_id1)
+    lock2 = AgencyConfigLockManager.get_lock(agency_id2)
+
+    assert lock1 is not lock2, "Different agencies should have different locks"
+
+
+def test_lock_for_same_agency():
+    agency_id = "agency1"
+
+    lock1 = AgencyConfigLockManager.get_lock(agency_id)
+    lock2 = AgencyConfigLockManager.get_lock(agency_id)
+
+    assert lock1 is lock2, "The same agency should return the same lock instance"
+
+
+def test_lock_concurrency_handling():
+    agency_id = "agency1"
+    lock = AgencyConfigLockManager.get_lock(agency_id)
+
+    # Define a shared resource
+    shared_resource = []
+
+    def task():
+        with lock:
+            shared_resource.append(1)
+
+    # Run tasks in parallel to simulate concurrent access
+    threads = [threading.Thread(target=task) for _ in range(10)]
+    for thread in threads:
+        thread.start()
+    for thread in threads:
+        thread.join()
+
+    assert len(shared_resource) == 10, "Concurrency issue: shared resource was accessed simultaneously"