UT: Add InMemoryFileRepository

monkey/tests/monkey_island/__init__.py

@@ -4,3 +4,4 @@ from .open_error_file_repository import OpenErrorFileRepository
 from .in_memory_agent_configuration_repository import InMemoryAgentConfigurationRepository
 from .in_memory_simulation_configuration import InMemorySimulationRepository
 from .in_memory_credentials_repository import InMemoryCredentialsRepository
+from .in_memory_file_repository import InMemoryFileRepository

monkey/tests/monkey_island/in_memory_file_repository.py

@@ -0,0 +1,33 @@
+import io
+import re
+from typing import BinaryIO, Dict
+
+from common.utils.code_utils import del_key
+from monkey_island.cc.repository import IFileRepository, UnknownRecordError
+
+
+class InMemoryFileRepository(IFileRepository):
+    def __init__(self):
+        self._files: Dict[str, bytes] = {}
+
+    def save_file(self, unsafe_file_name: str, file_contents: BinaryIO):
+        self._files[unsafe_file_name] = file_contents.read()
+
+    def open_file(self, unsafe_file_name: str) -> BinaryIO:
+        try:
+            return io.BytesIO(self._files[unsafe_file_name])
+        except KeyError:
+            raise UnknownRecordError(f"Unknown file {unsafe_file_name}")
+
+    def delete_file(self, unsafe_file_name: str):
+        del_key(self._files, "unsafe_file_name")
+
+    def delete_files_by_regex(self, file_name_regex: re.Pattern):
+        self._files = {
+            name: contents
+            for name, contents in self._files.items()
+            if not re.match(file_name_regex, name)
+        }
+
+    def delete_all_files(self):
+        self._files = {}