UT: Add delete_files_by_pattern() to SingleFileRepository and MockFileRepository

This commit is contained in:
Shreya Malviya 2022-09-28 16:10:30 +05:30
parent e1d139fde4
commit c1a4641ffe
2 changed files with 11 additions and 0 deletions

View File

@ -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

View File

@ -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("")