fix test

Author: minhkhul

nssp/delphi_nssp/patch.py

@@ -74,11 +74,12 @@ def good_patch_config(params, logger):
         logger.error("Custom flag is on, but patch section is missing.")
         valid_config = False
     else:
-        required_patch_keys = ["start_issue", "end_issue", "patch_dir", "source_dir", "source_host"]
+        required_patch_keys = ["start_issue", "end_issue", "patch_dir", "source_dir"]
 
         source_dir = params["patch"]["source_dir"]
         if not path.isdir(source_dir) or not listdir(source_dir):
             required_patch_keys.append("user")
+            required_patch_keys.append("source_host")
 
         missing_keys = [key for key in required_patch_keys if key not in patch_config]
         if missing_keys:

nssp/delphi_nssp/pull.py

@@ -51,29 +51,30 @@ def get_source_data(params, logger):
     remote_source_files = [f"{date.strftime('%Y%m%d')}.csv.gz" for date in dates]
 
     # Download source files
-    with ssh.open_sftp() as sftp:
+    sftp = ssh.open_sftp()
+    try:
+        sftp.stat(params["common"]["backup_dir"])
+    except IOError:
+        logger.error("Source backup directory does not exist on the remote server.")
+
+    sftp.chdir(params["common"]["backup_dir"])
+
+    num_files_transferred = 0
+    for remote_file_name in remote_source_files:
+        callback_for_filename = functools.partial(print_callback, remote_file_name, logger, progress_chunks=[0, 50])
+        local_file_path = path.join(params["patch"]["source_dir"], remote_file_name)
         try:
-            sftp.stat(params["common"]["backup_dir"])
+            sftp.stat(remote_file_name)
         except IOError:
-            logger.error("Source backup directory does not exist on the remote server.")
-
-        sftp.chdir(params["common"]["backup_dir"])
-
-        num_files_transferred = 0
-        for remote_file_name in remote_source_files:
-            callback_for_filename = functools.partial(print_callback, remote_file_name, logger, progress_chunks=[0, 50])
-            local_file_path = path.join(params["patch"]["source_dir"], remote_file_name)
-            try:
-                sftp.stat(remote_file_name)
-            except IOError:
-                logger.warning(
-                    "Source backup for this date does not exist on the remote server.",
-                    missing_filename=remote_file_name,
-                )
-                continue
-            sftp.get(remote_file_name, local_file_path, callback=callback_for_filename)
-            logger.info("Transfer finished", remote_file_name=remote_file_name, local_file_path=local_file_path)
-            num_files_transferred += 1
+            logger.warning(
+                "Source backup for this date does not exist on the remote server.",
+                missing_filename=remote_file_name,
+            )
+            continue
+        sftp.get(remote_file_name, local_file_path, callback=callback_for_filename)
+        logger.info("Transfer finished", remote_file_name=remote_file_name, local_file_path=local_file_path)
+        num_files_transferred += 1
+    ssh.close()
 
     if num_files_transferred == 0:
         logger.error("No source data was transferred. Check the source backup server for potential issues.")

nssp/tests/test_patch.py

@@ -99,7 +99,7 @@ def test_config_all_valid_configurations(self, mock_logger):
         mock_logger.info.assert_called_once_with("Good patch configuration.")
 
     @mock_patch('logging.Logger')
-    def test_config_user_param(self, mock_logger):
+    def test_config_user_and_host_param(self, mock_logger):
         # Case 6.1: pre-existing local source data,
         # so no "user" param in patch section needed
         patch_config = {
@@ -110,13 +110,12 @@ def test_config_user_param(self, mock_logger):
                 "patch_dir": "dir",
                 "start_issue": "2024-04-21",
                 "end_issue": "2024-04-22",
-                "source_dir": "./source_dir",
-                "source_host": "prod.server.edu"
+                "source_dir": "./source_dir"
             }
         }
         assert good_patch_config(patch_config, mock_logger)
 
-        # Case 6.2: source_dir does not exist, so "user" param is needed
+        # Case 6.2: source_dir does not exist, so "user" and "source_host" param is needed
         patch_config = {
             "common": {
                 "custom_run": True,
@@ -130,7 +129,7 @@ def test_config_user_param(self, mock_logger):
         }
         assert not good_patch_config(patch_config, mock_logger)
         mock_logger.error.assert_has_calls([
-            call("Patch section is missing required key(s)", missing_keys=["user"]),
+            call("Patch section is missing required key(s)", missing_keys=["user", "source_host"]),
         ])
 
     def test_get_patch_dates(self):