From bddee026fe76849c45c835cfc006b6ff2086de77 Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Wed, 28 Sep 2022 13:29:34 -0400
Subject: [PATCH] UT: Add InMemoryFileRepository

---
 monkey/tests/monkey_island/__init__.py        |  1 +
 .../in_memory_file_repository.py              | 33 +++++++++++++++++++
 2 files changed, 34 insertions(+)
 create mode 100644 monkey/tests/monkey_island/in_memory_file_repository.py

diff --git a/monkey/tests/monkey_island/__init__.py b/monkey/tests/monkey_island/__init__.py
index f7033842d..9dd1996f9 100644
--- a/monkey/tests/monkey_island/__init__.py
+++ b/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
diff --git a/monkey/tests/monkey_island/in_memory_file_repository.py b/monkey/tests/monkey_island/in_memory_file_repository.py
new file mode 100644
index 000000000..76778d226
--- /dev/null
+++ b/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 = {}