diff --git a/nalgonda/models/agency_config.py b/nalgonda/models/agency_config.py index ae1250f8..5cc98d56 100644 --- a/nalgonda/models/agency_config.py +++ b/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()) diff --git a/nalgonda/persistence/__init__.py b/nalgonda/persistence/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/nalgonda/persistence/agency_config_file_storage.py b/nalgonda/persistence/agency_config_file_storage.py new file mode 100644 index 00000000..c130f673 --- /dev/null +++ b/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") diff --git a/nalgonda/agency_config_lock_manager.py b/nalgonda/persistence/agency_config_lock_manager.py similarity index 100% rename from nalgonda/agency_config_lock_manager.py rename to nalgonda/persistence/agency_config_lock_manager.py diff --git a/nalgonda/persistence/agency_config_storage_interface.py b/nalgonda/persistence/agency_config_storage_interface.py new file mode 100644 index 00000000..e9e58462 --- /dev/null +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 97924be4..de38fde4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,6 +4,7 @@ version = "1.0.0" description = "Nalgonda is a FastAPI app to manage swarm agencies" authors = [ "Nikita Bobrovskiy <39348559+bonk1t@users.noreply.github.com>", + "Guilherme Parpinelli ", ] readme = "README.md" diff --git a/tests/persistence/test_agency_config_file_storage.py b/tests/persistence/test_agency_config_file_storage.py new file mode 100644 index 00000000..6d1dc182 --- /dev/null +++ b/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" diff --git a/tests/persistence/test_agency_config_lock_manager.py b/tests/persistence/test_agency_config_lock_manager.py new file mode 100644 index 00000000..64343154 --- /dev/null +++ b/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"