Island: Raise RemovalError from IMachineRepository.reset()

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
         """

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}")

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()