diff --git a/monkey/tests/monkey_island/mock_file_repository.py b/monkey/tests/monkey_island/mock_file_repository.py index 782c9838b..f828cdd5a 100644 --- a/monkey/tests/monkey_island/mock_file_repository.py +++ b/monkey/tests/monkey_island/mock_file_repository.py @@ -24,5 +24,8 @@ class MockFileRepository(IFileRepository): def delete_file(self, unsafe_file_name: str): pass + def delete_files_by_pattern(self, file_name_pattern: str): + pass + def delete_all_files(self): pass diff --git a/monkey/tests/monkey_island/single_file_repository.py b/monkey/tests/monkey_island/single_file_repository.py index 462969acb..66086f0f2 100644 --- a/monkey/tests/monkey_island/single_file_repository.py +++ b/monkey/tests/monkey_island/single_file_repository.py @@ -1,4 +1,5 @@ import io +import re from typing import BinaryIO from monkey_island.cc import repository @@ -8,9 +9,11 @@ from monkey_island.cc.repository import IFileRepository class SingleFileRepository(IFileRepository): def __init__(self): self._file = None + self._file_name = "" def save_file(self, unsafe_file_name: str, file_contents: BinaryIO): self._file = io.BytesIO(file_contents.read()) + self._file_name = unsafe_file_name def open_file(self, unsafe_file_name: str) -> BinaryIO: if self._file is None: @@ -19,6 +22,11 @@ class SingleFileRepository(IFileRepository): def delete_file(self, unsafe_file_name: str): self._file = None + self._file_name = "" + + def delete_files_by_pattern(self, file_name_pattern: str): + if re.match(file_name_pattern, self._file_name): + self.delete_file("") def delete_all_files(self): self.delete_file("")