Factor out constant and update tests

Author: krivard

changehc/delphi_changehc/constants.py

@@ -9,3 +9,5 @@
 NA = "NA"
 HRR = "hrr"
 FIPS = "fips"
+
+EXPECTED_FILES_PER_DROP = 7

changehc/delphi_changehc/download_ftp_files.py

@@ -8,6 +8,7 @@
 # third party
 import paramiko
 
+from .constants import EXPECTED_FILES_PER_DROP
 
 def print_callback(filename, bytes_so_far, bytes_total):
     """Log file transfer progress."""
@@ -32,8 +33,8 @@ def get_files_from_dir(sftp, filedate, out_path):
                 not path.exists(path.join(out_path, filename)):
             filepaths_to_download[filename] = path.join(out_path, filename)
 
-    # make sure we don't download more than 7 files per day
-    assert len(filepaths_to_download) <= 7, "more files dropped than expected"
+    # make sure we don't download too many files per day
+    assert len(filepaths_to_download) <= EXPECTED_FILES_PER_DROP, "more files dropped than expected"
 
     # download!
     for infile, outfile in filepaths_to_download.items():

changehc/tests/test_download_ftp_files.py

@@ -6,6 +6,7 @@
 
 # first party
 from delphi_changehc.download_ftp_files import *
+from delphi_changehc.constants import EXPECTED_FILES_PER_DROP
 
 class TestDownloadFTPFiles:
 
@@ -51,22 +52,18 @@ def test_get_files(self, mock_path):
         get_files_from_dir(one_new_one_old, "00005566", "")
         assert one_new_one_old.num_gets == 1
 
-        # When seven new files are present, AssertionError
-        new_file1 = self.FileAttr(dt.timestamp(dt.now()), "00001122_foo1")
-        new_file2 = self.FileAttr(dt.timestamp(dt.now()), "00001122_foo2")
-        new_file3 = self.FileAttr(dt.timestamp(dt.now()), "00001122_foo3")
-        new_file4 = self.FileAttr(dt.timestamp(dt.now()), "00001122_foo4")
-        new_file5 = self.FileAttr(dt.timestamp(dt.now()), "00001122_foo5")
-        new_file6 = self.FileAttr(dt.timestamp(dt.now()), "00001122_foo6")
-        new_file7 = self.FileAttr(dt.timestamp(dt.now()), "00001122_foo7")
-        seven_new = self.MockSFTP([new_file1, new_file2, new_file3, new_file4,
-                                    new_file5, new_file6, new_file7])
+        # When too many new files are present, AssertionError
+        file_batch = [
+            self.FileAttr(dt.timestamp(dt.now()), f"00001122_foo{i}")
+            for i in range(EXPECTED_FILES_PER_DROP + 1)
+        ]
+        too_many_new = self.MockSFTP(file_batch)
         with pytest.raises(AssertionError):
-            get_files_from_dir(seven_new, "00001122", "")
+            get_files_from_dir(too_many_new, "00001122", "")
 
         # When the file already exists, no files are downloaded
         mock_path.exists.return_value = True
-        one_exists = self.MockSFTP([new_file1])
+        one_exists = self.MockSFTP([file_batch[0]])
         get_files_from_dir(one_new, "00001122", "")
         assert one_exists.num_gets == 0