p15670423
/
monkey
1
0
Fork 5
Code Issues Pull Requests 3 Projects Releases Wiki Activity

tests: Fix ransomware tests and move tests for get_all_files_in_directory from ransomware/test_utils.py to utils/test_dir_utils.py

This commit is contained in:
Shreya 2021-06-22 20:07:38 +05:30
parent a2ebe3386f
commit efef40edf9
2 changed files with 46 additions and 48 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

48
monkey/tests/unit_tests/infection_monkey/ransomware/test_utils.py
View File

@ -1,5 +1,3 @@
import os
import infection_monkey.ransomware.utils
VALID_FILE_EXTENSION_1 = "file.3ds"
@ -43,49 +41,3 @@ def test_get_files_to_encrypt__valid_files(monkeypatch):
expected_return_value = [VALID_FILE_EXTENSION_1, VALID_FILE_EXTENSION_2]
assert infection_monkey.ransomware.utils.get_files_to_encrypt("") == expected_return_value
def test_get_all_files_in_directory__no_files(tmpdir, monkeypatch):
subdir1 = os.path.join(tmpdir, SUBDIR_1)
subdir2 = os.path.join(tmpdir, SUBDIR_2)
subdirs = [subdir1, subdir2]
for subdir in subdirs:
os.mkdir(subdir)
all_items_in_dir = subdirs
monkeypatch.setattr("os.listdir", lambda _: all_items_in_dir)
expected_return_value = []
assert (
infection_monkey.ransomware.utils.get_all_files_in_directory(tmpdir)
== expected_return_value
)
def test_get_all_files_in_directory__has_files(tmpdir, monkeypatch):
subdir1 = os.path.join(tmpdir, SUBDIR_1)
subdir2 = os.path.join(tmpdir, SUBDIR_2)
subdirs = [subdir1, subdir2]
file1 = os.path.join(tmpdir, VALID_FILE_EXTENSION_1)
file2 = os.path.join(tmpdir, INVALID_FILE_EXTENSION_1)
file3 = os.path.join(tmpdir, VALID_FILE_EXTENSION_2)
file4 = os.path.join(tmpdir, INVALID_FILE_EXTENSION_2)
files = [file1, file2, file3, file4]
for subdir in subdirs:
os.mkdir(subdir)
for file in files:
with open(file, "w") as _:
pass
all_items_in_dir = subdirs + files
monkeypatch.setattr("os.listdir", lambda _: all_items_in_dir)
expected_return_value = files
assert (
infection_monkey.ransomware.utils.get_all_files_in_directory(tmpdir)
== expected_return_value
)

46
monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py Normal file
View File

@ -0,0 +1,46 @@
import os
from infection_monkey.utils.dir_utils import get_all_files_in_directory
FILE_1 = "file.jpg.zip"
FILE_2 = "file.xyz"
SUBDIR_1 = "subdir1"
SUBDIR_2 = "subdir2"
def test_get_all_files_in_directory__no_files(tmpdir, monkeypatch):
subdir1 = os.path.join(tmpdir, SUBDIR_1)
subdir2 = os.path.join(tmpdir, SUBDIR_2)
subdirs = [subdir1, subdir2]
for subdir in subdirs:
os.mkdir(subdir)
all_items_in_dir = subdirs
monkeypatch.setattr("os.listdir", lambda _: all_items_in_dir)
expected_return_value = []
assert get_all_files_in_directory(tmpdir) == expected_return_value
def test_get_all_files_in_directory__has_files(tmpdir, monkeypatch):
subdir1 = os.path.join(tmpdir, SUBDIR_1)
subdir2 = os.path.join(tmpdir, SUBDIR_2)
subdirs = [subdir1, subdir2]
file1 = os.path.join(tmpdir, FILE_1)
file2 = os.path.join(tmpdir, FILE_2)
files = [file1, file2]
for subdir in subdirs:
os.mkdir(subdir)
for file in files:
with open(file, "w") as _:
pass
all_items_in_dir = subdirs + files
monkeypatch.setattr("os.listdir", lambda _: all_items_in_dir)
expected_return_value = files
assert get_all_files_in_directory(tmpdir) == expected_return_value