From efef40edf911573a3958e54a9620a580619293fe Mon Sep 17 00:00:00 2001
From: Shreya <shreya.malviya@gmail.com>
Date: Tue, 22 Jun 2021 20:07:38 +0530
Subject: [PATCH] tests: Fix ransomware tests and move tests for
 get_all_files_in_directory from ransomware/test_utils.py to
 utils/test_dir_utils.py

---
 .../infection_monkey/ransomware/test_utils.py | 48 -------------------
 .../infection_monkey/utils/test_dir_utils.py  | 46 ++++++++++++++++++
 2 files changed, 46 insertions(+), 48 deletions(-)
 create mode 100644 monkey/tests/unit_tests/infection_monkey/utils/test_dir_utils.py

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