Island: Raise RemovalError from IMachineRepository.reset()

This commit is contained in:
Mike Salvatore 2022-08-30 07:34:20 -04:00
parent 383cfdfefe
commit bf5e54ebc9
3 changed files with 15 additions and 2 deletions

View File

@ -68,4 +68,7 @@ class IMachineRepository(ABC):
def reset(self): def reset(self):
""" """
Removes all data from the repository Removes all data from the repository
:raises RemovalError: If an error occurred while attempting to remove all `Machines` from
the repository
""" """

View File

@ -7,7 +7,7 @@ from pymongo import MongoClient
from common.types import HardwareID from common.types import HardwareID
from monkey_island.cc.models import Machine, MachineID 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 from .consts import MONGO_OBJECT_ID_KEY
@ -96,4 +96,7 @@ class MongoMachineRepository(IMachineRepository):
return Machine(**mongo_record) return Machine(**mongo_record)
def reset(self): def reset(self):
self._machines_collection.drop() try:
self._machines_collection.drop()
except Exception as err:
raise RemovalError(f"Error resetting the repository: {err}")

View File

@ -10,6 +10,7 @@ from monkey_island.cc.models import Machine
from monkey_island.cc.repository import ( from monkey_island.cc.repository import (
IMachineRepository, IMachineRepository,
MongoMachineRepository, MongoMachineRepository,
RemovalError,
RetrievalError, RetrievalError,
StorageError, StorageError,
UnknownRecordError, UnknownRecordError,
@ -77,6 +78,7 @@ def error_raising_mock_mongo_client() -> mongomock.MongoClient:
mongo_client.monkey_island.machines.replace_one = MagicMock( mongo_client.monkey_island.machines.replace_one = MagicMock(
side_effect=Exception("some exception") side_effect=Exception("some exception")
) )
mongo_client.monkey_island.machines.drop = MagicMock(side_effect=Exception("some exception"))
return mongo_client return mongo_client
@ -236,3 +238,8 @@ def test_usable_after_reset(machine_repository):
new_machine = machine_repository.create_machine() new_machine = machine_repository.create_machine()
assert new_machine == machine_repository.get_machine_by_id(new_machine.id) 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()