Move CSV-generation code into its own function and write unit tests

Author: gitname

main.py

@@ -75,6 +75,25 @@ def get_file_paths_from_csv_file(path_to_csv_file):
         return file_paths
 
 
+def write_to_csv_file(file_paths, path_to_output_file):
+    """Writes the specified file paths--and some of their constituent elements--to the specified output file."""
+
+    with open(path_to_output_file, 'w', encoding='utf-8', newline='') as f:
+        writer = csv.writer(f)
+
+        # Populate the header row (i.e. column names).
+        header_row = ["File", "Folder", "Filename", "Suffix (only)"]
+        writer.writerow(header_row)
+
+        # Populate the data rows with the file path, folder path (no filename), filename, and file suffix.
+        data_rows = []
+        for file_path in file_paths:
+            file_suffix = pathlib.Path(file_path).suffix
+            (folder_path, filename) = os.path.split(file_path)
+            data_rows.append([file_path, folder_path, filename, file_suffix])
+        writer.writerows(data_rows)
+
+
 def main(path_to_folder, path_to_csv_file, path_to_output_file=None, log_level=logging.NOTSET):
     """Entrypoint to the script."""
 
@@ -107,20 +126,7 @@ def main(path_to_folder, path_to_csv_file, path_to_output_file=None, log_level=l
 
     # Either write the results to a CSV file or print them to STDOUT.
     if type(path_to_output_file) is str:
-        with open(path_to_output_file, 'w', encoding='utf-8', newline='') as f:
-            writer = csv.writer(f)
-
-            # Populate the header row (i.e. column names).
-            header_row = ["File", "Folder", "Filename", "Suffix (only)"]
-            writer.writerow(header_row)
-
-            # Populate the data rows with the file path, folder path (no filename), filename, and file suffix.
-            data_rows = []
-            for path in sorted_differences:
-                file_suffix = pathlib.Path(path).suffix
-                (folder_path, filename) = os.path.split(path)
-                data_rows.append([path, folder_path, filename, file_suffix])
-            writer.writerows(data_rows)
+        write_to_csv_file(sorted_differences, path_to_output_file)
     else:
         print('\n'.join(sorted_differences))
 

test/test.py

@@ -1,7 +1,13 @@
+import os
 import unittest
-from main import get_file_paths_from_csv_file, get_file_paths_within_folder
+from main import (
+    get_file_paths_from_csv_file,
+    get_file_paths_within_folder,
+    write_to_csv_file
+)
 from os import path
 import tempfile
+import csv
 
 
 class TestGetFilePathsFromCsvFile(unittest.TestCase):
@@ -65,5 +71,50 @@ def test_child_folder(self):
         self.assertEqual(paths, expectation)
 
 
+class TestWriteToCsvFile(unittest.TestCase):
+    header_row = ["File", "Folder", "Filename", "Suffix (only)"]
+    temp_file = None
+
+    def setUp(self):
+        """Create a temporary file before each test."""
+        self.temp_file = tempfile.NamedTemporaryFile(suffix='.csv', delete=False)
+
+    def tearDown(self):
+        """Close and delete the temporary file after each test."""
+        self.temp_file.close()
+        os.remove(self.temp_file.name)
+
+    def test_empty_list(self):
+        write_to_csv_file([], self.temp_file.name)
+        with open(self.temp_file.name, 'r', encoding='utf8', newline='') as csv_file:
+            reader = csv.reader(csv_file)
+            rows = list(reader)
+            self.assertEqual(rows, [self.header_row])
+
+    def test_nonempty_list(self):
+        file_paths = [
+            r"C:/A/B/C",  # forward slashes
+            r"C:\A\B\C",  # no suffix
+            r"C:\A\B\C.d",  # suffix
+            r"C:\A\B\C.d.e",  # multiple suffixes
+            r"C:\A\B\C d.e",  # space in file name
+            r"C:\A B\C",  # space in folder name
+        ]
+        expected_rows = [
+            self.header_row,
+            [r"C:/A/B/C", r"C:/A/B", r"C", r""],
+            [r"C:\A\B\C", r"C:\A\B", r"C", r""],
+            [r"C:\A\B\C.d", r"C:\A\B", r"C.d", r".d"],
+            [r"C:\A\B\C.d.e", r"C:\A\B", r"C.d.e", r".e"],
+            [r"C:\A\B\C d.e", r"C:\A\B", r"C d.e", r".e"],
+            [r"C:\A B\C", r"C:\A B", r"C", r""],
+        ]
+        write_to_csv_file(file_paths, self.temp_file.name)
+        with open(self.temp_file.name, 'r', encoding='utf8', newline='') as csv_file:
+            reader = csv.reader(csv_file)
+            rows = list(reader)
+            self.assertEqual(rows, expected_rows)
+
+
 if __name__ == '__main__':
     unittest.main()