forked from p15670423/monkey
Island: Raise RemovalError from IMachineRepository.reset()
This commit is contained in:
parent
383cfdfefe
commit
bf5e54ebc9
|
@ -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
|
||||||
"""
|
"""
|
||||||
|
|
|
@ -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}")
|
||||||
|
|
|
@ -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()
|
||||||
|
|
Loading…
Reference in New Issue