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:
parent
a2ebe3386f
commit
efef40edf9
|
@ -1,5 +1,3 @@
|
||||||
import os
|
|
||||||
|
|
||||||
import infection_monkey.ransomware.utils
|
import infection_monkey.ransomware.utils
|
||||||
|
|
||||||
VALID_FILE_EXTENSION_1 = "file.3ds"
|
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]
|
expected_return_value = [VALID_FILE_EXTENSION_1, VALID_FILE_EXTENSION_2]
|
||||||
assert infection_monkey.ransomware.utils.get_files_to_encrypt("") == expected_return_value
|
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
|
|
||||||
)
|
|
||||||
|
|
|
@ -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
|
Loading…
Reference in New Issue