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):
|
||||
"""
|
||||
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 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):
|
||||
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 (
|
||||
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()
|
||||
|
|
Loading…
Reference in New Issue