From e2dfd6a5e3c7d0028985a990982fc040371c6b17 Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Tue, 22 Jun 2021 12:15:03 -0400
Subject: [PATCH] agent: Simplify get_all_files_in_directory() with list
 comprehension

---
 monkey/infection_monkey/utils/dir_utils.py    |  8 +--
 .../infection_monkey/utils/test_dir_utils.py  | 50 +++++++++++--------
 2 files changed, 34 insertions(+), 24 deletions(-)

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