diff --git a/monkey/tests/common/SingleFileRepository.py b/monkey/tests/common/SingleFileRepository.py new file mode 100644 index 000000000..d00ee9ae9 --- /dev/null +++ b/monkey/tests/common/SingleFileRepository.py @@ -0,0 +1,23 @@ +import io +from typing import BinaryIO + +from monkey_island.cc.repository import FileRetrievalError, IFileRepository + + +class SingleFileRepository(IFileRepository): + def __init__(self): + self._file = None + + def save_file(self, unsafe_file_name: str, file_contents: BinaryIO): + self._file = io.BytesIO(file_contents.read()) + + def open_file(self, unsafe_file_name: str) -> BinaryIO: + if self._file is None: + raise FileRetrievalError() + return self._file + + def delete_file(self, unsafe_file_name: str): + self._file = None + + def delete_all_files(self): + self.delete_file("") diff --git a/monkey/tests/monkey_island/__init__.py b/monkey/tests/monkey_island/__init__.py new file mode 100644 index 000000000..7bd9a314d --- /dev/null +++ b/monkey/tests/monkey_island/__init__.py @@ -0,0 +1 @@ +from .single_file_repository import SingleFileRepository diff --git a/monkey/tests/monkey_island/single_file_repository.py b/monkey/tests/monkey_island/single_file_repository.py new file mode 100644 index 000000000..d00ee9ae9 --- /dev/null +++ b/monkey/tests/monkey_island/single_file_repository.py @@ -0,0 +1,23 @@ +import io +from typing import BinaryIO + +from monkey_island.cc.repository import FileRetrievalError, IFileRepository + + +class SingleFileRepository(IFileRepository): + def __init__(self): + self._file = None + + def save_file(self, unsafe_file_name: str, file_contents: BinaryIO): + self._file = io.BytesIO(file_contents.read()) + + def open_file(self, unsafe_file_name: str) -> BinaryIO: + if self._file is None: + raise FileRetrievalError() + return self._file + + def delete_file(self, unsafe_file_name: str): + self._file = None + + def delete_all_files(self): + self.delete_file("") diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py index 2642c0758..50da88915 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_upload.py @@ -1,12 +1,10 @@ -import io -from typing import BinaryIO - import pytest from tests.common import StubDIContainer +from tests.monkey_island import SingleFileRepository from tests.unit_tests.monkey_island.conftest import get_url_for_resource from tests.utils import raise_ -from monkey_island.cc.repository import FileRetrievalError, IFileRepository +from monkey_island.cc.repository import IFileRepository from monkey_island.cc.resources.pba_file_upload import LINUX_PBA_TYPE, WINDOWS_PBA_TYPE, FileUpload TEST_FILE_CONTENTS = b"m0nk3y" @@ -40,28 +38,9 @@ def mock_get_config_value(monkeypatch): ) -class MockFileRepository(IFileRepository): - def __init__(self): - self._file = None - - def save_file(self, unsafe_file_name: str, file_contents: BinaryIO): - self._file = io.BytesIO(file_contents.read()) - - def open_file(self, unsafe_file_name: str) -> BinaryIO: - if self._file is None: - raise FileRetrievalError() - return self._file - - def delete_file(self, unsafe_file_name: str): - self._file = None - - def delete_all_files(self): - self.delete_file("") - - @pytest.fixture def file_repository(): - return MockFileRepository() + return SingleFileRepository() @pytest.fixture