diff --git a/monkey/infection_monkey/utils/dir_utils.py b/monkey/infection_monkey/utils/dir_utils.py index 4350a279d..b64e4e13f 100644 --- a/monkey/infection_monkey/utils/dir_utils.py +++ b/monkey/infection_monkey/utils/dir_utils.py @@ -1,6 +1,10 @@ from pathlib import Path -from typing import List +from typing import Callable, List def get_all_files_in_directory(dir_path: Path) -> List[Path]: return [f for f in dir_path.iterdir() if f.is_file()] + + +def filter_files(files: List[Path], file_filter: Callable[[Path], bool]): + return [f for f in files if file_filter(f)] 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 7369c5f41..604efc9cb 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,7 @@ -from infection_monkey.utils.dir_utils import get_all_files_in_directory +from infection_monkey.utils.dir_utils import ( + filter_files, + get_all_files_in_directory, +) FILE_1 = "file.jpg.zip" FILE_2 = "file.xyz" @@ -51,3 +54,22 @@ def test_get_all_files_in_directory__subdir_has_files(tmp_path, monkeypatch): expected_return_value = sorted(files) assert sorted(get_all_files_in_directory(tmp_path)) == expected_return_value + + +def test_filter_files__no_results(tmp_path): + add_files_to_dir(tmp_path) + + files_in_dir = get_all_files_in_directory(tmp_path) + filtered_files = filter_files(files_in_dir, lambda _: False) + + assert len(filtered_files) == 0 + + +def test_filter_files__all_true(tmp_path): + files = add_files_to_dir(tmp_path) + expected_return_value = sorted(files) + + files_in_dir = get_all_files_in_directory(tmp_path) + filtered_files = filter_files(files_in_dir, lambda _: True) + + assert sorted(filtered_files) == expected_return_value