diff --git a/monkey/tests/unit_tests/infection_monkey/ransomware/test_utils.py b/monkey/tests/unit_tests/infection_monkey/ransomware/test_utils.py index f3c0bea39..37a7f059d 100644 --- a/monkey/tests/unit_tests/infection_monkey/ransomware/test_utils.py +++ b/monkey/tests/unit_tests/infection_monkey/ransomware/test_utils.py @@ -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 - ) diff --git a/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py b/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py new file mode 100644 index 000000000..22b2ba0c9 --- /dev/null +++ b/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py @@ -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