From bf5e54ebc9b164e7b389ab4d002514c8aaa4fd13 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 30 Aug 2022 07:34:20 -0400 Subject: [PATCH] Island: Raise RemovalError from IMachineRepository.reset() --- monkey/monkey_island/cc/repository/i_machine_repository.py | 3 +++ .../cc/repository/mongo_machine_repository.py | 7 +++++-- .../cc/repository/test_mongo_machine_repository.py | 7 +++++++ 3 files changed, 15 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/repository/i_machine_repository.py b/monkey/monkey_island/cc/repository/i_machine_repository.py index 6d6832669..568564009 100644 --- a/monkey/monkey_island/cc/repository/i_machine_repository.py +++ b/monkey/monkey_island/cc/repository/i_machine_repository.py @@ -68,4 +68,7 @@ class IMachineRepository(ABC): def reset(self): """ Removes all data from the repository + + :raises RemovalError: If an error occurred while attempting to remove all `Machines` from + the repository """ diff --git a/monkey/monkey_island/cc/repository/mongo_machine_repository.py b/monkey/monkey_island/cc/repository/mongo_machine_repository.py index e3d59169a..318153e51 100644 --- a/monkey/monkey_island/cc/repository/mongo_machine_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_machine_repository.py @@ -7,7 +7,7 @@ from pymongo import MongoClient from common.types import HardwareID from monkey_island.cc.models import Machine, MachineID -from . import IMachineRepository, RetrievalError, StorageError, UnknownRecordError +from . import IMachineRepository, RemovalError, RetrievalError, StorageError, UnknownRecordError from .consts import MONGO_OBJECT_ID_KEY @@ -96,4 +96,7 @@ class MongoMachineRepository(IMachineRepository): return Machine(**mongo_record) def reset(self): - self._machines_collection.drop() + try: + self._machines_collection.drop() + except Exception as err: + raise RemovalError(f"Error resetting the repository: {err}") diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_machine_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_machine_repository.py index 9ff00c64b..7851f8e20 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_machine_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_machine_repository.py @@ -10,6 +10,7 @@ from monkey_island.cc.models import Machine from monkey_island.cc.repository import ( IMachineRepository, MongoMachineRepository, + RemovalError, RetrievalError, StorageError, UnknownRecordError, @@ -77,6 +78,7 @@ def error_raising_mock_mongo_client() -> mongomock.MongoClient: mongo_client.monkey_island.machines.replace_one = MagicMock( side_effect=Exception("some exception") ) + mongo_client.monkey_island.machines.drop = MagicMock(side_effect=Exception("some exception")) return mongo_client @@ -236,3 +238,8 @@ def test_usable_after_reset(machine_repository): new_machine = machine_repository.create_machine() assert new_machine == machine_repository.get_machine_by_id(new_machine.id) + + +def test_reset__removal_error(error_raising_machine_repository): + with pytest.raises(RemovalError): + error_raising_machine_repository.reset()