diff --git a/monkey/monkey_island/cc/repository/file_agent_configuration_repository.py b/monkey/monkey_island/cc/repository/file_agent_configuration_repository.py index fcccd49a1..89fc57960 100644 --- a/monkey/monkey_island/cc/repository/file_agent_configuration_repository.py +++ b/monkey/monkey_island/cc/repository/file_agent_configuration_repository.py @@ -5,6 +5,7 @@ from common.configuration import ( AgentConfiguration, AgentConfigurationSchema, ) +from monkey_island.cc import repository from monkey_island.cc.repository import ( IAgentConfigurationRepository, IFileRepository, @@ -25,10 +26,10 @@ class FileAgentConfigurationRepository(IAgentConfigurationRepository): configuration_json = f.read().decode() return self._schema.loads(configuration_json) - # TODO: Handle FileRetrievalError vs FileNotFoundError - # https://github.com/guardicore/monkey/blob/e8001d8cf76340e42bf17ff62523bd2d85fc4841/monkey/monkey_island/cc/repository/file_storage/local_storage_file_repository.py#L47-L50 - except RetrievalError: + except repository.FileNotFoundError: return self._schema.loads(DEFAULT_AGENT_CONFIGURATION) + except Exception as err: + raise RetrievalError(f"Error retrieving the agent configuration: {err}") def store_configuration(self, agent_configuration: AgentConfiguration): configuration_json = self._schema.dumps(agent_configuration) diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_configuration_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_configuration_repository.py index 7ad066623..93a80dea6 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_configuration_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_configuration_repository.py @@ -1,8 +1,9 @@ +import pytest from tests.common.example_agent_configuration import AGENT_CONFIGURATION -from tests.monkey_island import SingleFileRepository +from tests.monkey_island import OpenErrorFileRepository, SingleFileRepository from common.configuration import DEFAULT_AGENT_CONFIGURATION, AgentConfigurationSchema -from monkey_island.cc.repository import FileAgentConfigurationRepository +from monkey_island.cc.repository import FileAgentConfigurationRepository, RetrievalError def test_store_agent_config(): @@ -24,3 +25,10 @@ def test_get_default_agent_config(): retrieved_agent_configuration = repository.get_configuration() assert retrieved_agent_configuration == default_agent_configuration + + +def test_get_agent_config_retrieval_error(): + repository = FileAgentConfigurationRepository(OpenErrorFileRepository()) + + with pytest.raises(RetrievalError): + repository.get_configuration()