diff --git a/monkey/infection_monkey/utils/dir_utils.py b/monkey/infection_monkey/utils/dir_utils.py index 5562dd6ec..62f592c5e 100644 --- a/monkey/infection_monkey/utils/dir_utils.py +++ b/monkey/infection_monkey/utils/dir_utils.py @@ -1,8 +1,8 @@ -import os +from pathlib import Path from typing import List def get_all_files_in_directory(dir_path: str) -> List: - return list( - filter(os.path.isfile, [os.path.join(dir_path, item) for item in os.listdir(dir_path)]) - ) + path = Path(dir_path) + + return [str(f) for f in path.iterdir() if f.is_file()] 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 index 22b2ba0c9..7156662ba 100644 --- a/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py +++ b/monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py @@ -1,4 +1,5 @@ import os +from pathlib import Path from infection_monkey.utils.dir_utils import get_all_files_in_directory @@ -8,39 +9,48 @@ 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) +def add_subdirs_to_dir(parent_dir): + subdir1 = os.path.join(parent_dir, SUBDIR_1) + subdir2 = os.path.join(parent_dir, 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) + return subdirs + + +def add_files_to_dir(parent_dir): + file1 = os.path.join(parent_dir, FILE_1) + file2 = os.path.join(parent_dir, FILE_2) + files = [file1, file2] + + for f in files: + Path(f).touch() + + return files + + +def test_get_all_files_in_directory__no_files(tmpdir, monkeypatch): + add_subdirs_to_dir(tmpdir) 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] + add_subdirs_to_dir(tmpdir) + files = add_files_to_dir(tmpdir) - file1 = os.path.join(tmpdir, FILE_1) - file2 = os.path.join(tmpdir, FILE_2) - files = [file1, file2] + expected_return_value = sorted(files) + assert sorted(get_all_files_in_directory(tmpdir)) == expected_return_value - for subdir in subdirs: - os.mkdir(subdir) - for file in files: - with open(file, "w") as _: - pass +def test_get_all_files_in_directory__subdir_has_files(tmpdir, monkeypatch): + subdirs = add_subdirs_to_dir(tmpdir) + add_files_to_dir(subdirs[0]) - all_items_in_dir = subdirs + files - monkeypatch.setattr("os.listdir", lambda _: all_items_in_dir) + files = add_files_to_dir(tmpdir) - expected_return_value = files - assert get_all_files_in_directory(tmpdir) == expected_return_value + expected_return_value = sorted(files) + assert sorted(get_all_files_in_directory(tmpdir)) == expected_return_value