diff --git a/monkey/tests/monkey_island/__init__.py b/monkey/tests/monkey_island/__init__.py index 7bd9a314d..aed4d07f3 100644 --- a/monkey/tests/monkey_island/__init__.py +++ b/monkey/tests/monkey_island/__init__.py @@ -1 +1,3 @@ from .single_file_repository import SingleFileRepository +from .mock_file_repository import MockFileRepository, FILE_CONTENTS, FILE_NAME +from .open_error_file_repository import OpenErrorFileRepository diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/mock_file_repository.py b/monkey/tests/monkey_island/mock_file_repository.py similarity index 100% rename from monkey/tests/unit_tests/monkey_island/cc/resources/mock_file_repository.py rename to monkey/tests/monkey_island/mock_file_repository.py diff --git a/monkey/tests/monkey_island/open_error_file_repository.py b/monkey/tests/monkey_island/open_error_file_repository.py new file mode 100644 index 000000000..c13559613 --- /dev/null +++ b/monkey/tests/monkey_island/open_error_file_repository.py @@ -0,0 +1,10 @@ +from typing import BinaryIO + +from monkey_island.cc.repository import RetrievalError + +from . import MockFileRepository + + +class OpenErrorFileRepository(MockFileRepository): + def open_file(self, unsafe_file_name: str) -> BinaryIO: + raise RetrievalError("Error retrieving file") diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py b/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py index 532282e8b..26fe24821 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/conftest.py @@ -1,17 +1,15 @@ -from typing import BinaryIO from unittest.mock import MagicMock import flask_jwt_extended import pytest from tests.common import StubDIContainer +from tests.monkey_island import OpenErrorFileRepository from tests.unit_tests.monkey_island.conftest import init_mock_app import monkey_island.cc.app import monkey_island.cc.resources.auth.auth import monkey_island.cc.resources.island_mode -from monkey_island.cc.repository import IFileRepository, RetrievalError - -from .mock_file_repository import MockFileRepository +from monkey_island.cc.repository import IFileRepository @pytest.fixture @@ -45,11 +43,6 @@ def get_mock_app(container): return app -class OpenErrorFileRepository(MockFileRepository): - def open_file(self, unsafe_file_name: str) -> BinaryIO: - raise RetrievalError("Error retrieving file") - - @pytest.fixture def open_error_flask_client(build_flask_client): container = StubDIContainer() diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py index eb189b4e7..31ae0309a 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_pba_file_download.py @@ -1,12 +1,11 @@ import pytest from tests.common import StubDIContainer +from tests.monkey_island import FILE_CONTENTS, FILE_NAME, MockFileRepository from tests.unit_tests.monkey_island.conftest import get_url_for_resource from monkey_island.cc.repository import IFileRepository from monkey_island.cc.resources.pba_file_download import PBAFileDownload -from .mock_file_repository import FILE_CONTENTS, FILE_NAME, MockFileRepository - @pytest.fixture def flask_client(build_flask_client):