forked from p34709852/monkey
UT: Extract SingleFileRepository from test_pba_file_upload.py
This commit is contained in:
parent
91476a7a06
commit
ace3eb8718
|
@ -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("")
|
|
@ -0,0 +1 @@
|
||||||
|
from .single_file_repository import SingleFileRepository
|
|
@ -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("")
|
|
@ -1,12 +1,10 @@
|
||||||
import io
|
|
||||||
from typing import BinaryIO
|
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
from tests.common import StubDIContainer
|
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.unit_tests.monkey_island.conftest import get_url_for_resource
|
||||||
from tests.utils import raise_
|
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
|
from monkey_island.cc.resources.pba_file_upload import LINUX_PBA_TYPE, WINDOWS_PBA_TYPE, FileUpload
|
||||||
|
|
||||||
TEST_FILE_CONTENTS = b"m0nk3y"
|
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
|
@pytest.fixture
|
||||||
def file_repository():
|
def file_repository():
|
||||||
return MockFileRepository()
|
return SingleFileRepository()
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
|
|
Loading…
Reference in New Issue