From d49e1708e1069371819cd07f51542db61f2e2edb Mon Sep 17 00:00:00 2001
From: Adam McDonagh <27215476+adammcdonagh@users.noreply.github.com>
Date: Tue, 31 Oct 2023 10:46:10 +0000
Subject: [PATCH] Add additional tests and fix return code when no exception

---
 src/opentaskpy/cli/task_run.py                |  5 ++-
 .../batch-basic-invalid-execution-host.json   | 10 ++++++
 test/cfg/executions/df-invalid-host.json      | 12 +++++++
 tests/test_task_run.py                        | 12 +++++++
 tests/test_taskhandler_transfer_sftp.py       | 31 +++++++++++++++++++
 5 files changed, 69 insertions(+), 1 deletion(-)
 create mode 100644 test/cfg/batch/batch-basic-invalid-execution-host.json
 create mode 100644 test/cfg/executions/df-invalid-host.json

diff --git a/src/opentaskpy/cli/task_run.py b/src/opentaskpy/cli/task_run.py
index 04953b9..c4475b0 100644
--- a/src/opentaskpy/cli/task_run.py
+++ b/src/opentaskpy/cli/task_run.py
@@ -118,13 +118,16 @@ def main() -> None:
     task_run_obj = taskrun.TaskRun(args.taskId, CONFIG_PATH, noop=args.noop)
 
     try:
-        task_run_obj.run()
+        result = task_run_obj.run()
     except Exception as ex:  # pylint: disable=broad-exception-caught
         logger.error(f"Error running task: {ex}")
         if logger.getEffectiveLevel() <= 12:
             raise ex
         sys.exit(1)
 
+    if not result:
+        sys.exit(1)
+
 
 if __name__ == "__main__":
     main()
diff --git a/test/cfg/batch/batch-basic-invalid-execution-host.json b/test/cfg/batch/batch-basic-invalid-execution-host.json
new file mode 100644
index 0000000..3117434
--- /dev/null
+++ b/test/cfg/batch/batch-basic-invalid-execution-host.json
@@ -0,0 +1,10 @@
+{
+  "type": "batch",
+  "tasks": [
+    {
+      "order_id": 1,
+      "task_id": "df-invalid-host",
+      "timeout": 60
+    }
+  ]
+}
diff --git a/test/cfg/executions/df-invalid-host.json b/test/cfg/executions/df-invalid-host.json
new file mode 100644
index 0000000..36099b7
--- /dev/null
+++ b/test/cfg/executions/df-invalid-host.json
@@ -0,0 +1,12 @@
+{
+  "type": "execution",
+  "hosts": ["123"],
+  "directory": "/tmp",
+  "command": "df -h",
+  "protocol": {
+    "name": "ssh",
+    "credentials": {
+      "username": "{{ SSH_USERNAME }}"
+    }
+  }
+}
diff --git a/tests/test_task_run.py b/tests/test_task_run.py
index 61ea7b9..08c0235 100644
--- a/tests/test_task_run.py
+++ b/tests/test_task_run.py
@@ -92,6 +92,12 @@ def test_execution_basic_binary(env_vars, setup_ssh_keys, root_dir):
     assert run_task_run("df")["returncode"] == 0
 
 
+def test_execution_invalid_host(env_vars, setup_ssh_keys, root_dir):
+    # Use the "binary" to trigger the job with command line arguments
+
+    assert run_task_run("df-invalid-host")["returncode"] == 1
+
+
 def test_batch_basic_binary(env_vars, setup_ssh_keys, root_dir):
     # Use the "binary" to trigger the job with command line arguments
 
@@ -103,6 +109,12 @@ def test_batch_basic_binary(env_vars, setup_ssh_keys, root_dir):
     assert run_task_run("batch-basic")["returncode"] == 0
 
 
+def test_batch_execution_invalid_host(env_vars, setup_ssh_keys, root_dir):
+    # Use the "binary" to trigger the job with command line arguments
+
+    assert run_task_run("batch-basic-invalid-execution-host")["returncode"] == 1
+
+
 def test_binary_invalid_config_file(env_vars, setup_ssh_keys, root_dir):
     # Use the "binary" to trigger the job with command line arguments
 
diff --git a/tests/test_taskhandler_transfer_sftp.py b/tests/test_taskhandler_transfer_sftp.py
index b4eff6d..88d5f51 100644
--- a/tests/test_taskhandler_transfer_sftp.py
+++ b/tests/test_taskhandler_transfer_sftp.py
@@ -30,6 +30,24 @@
     ],
 }
 
+# Non existent source directory
+sftp_task_non_existent_source_dir_definition = {
+    "type": "transfer",
+    "source": {
+        "hostname": "172.16.0.21",
+        "directory": "/home/application/testFiles/src/nonexistentdir",
+        "fileRegex": ".*taskhandler.*\\.txt",
+        "protocol": {"name": "sftp", "credentials": {"username": "application"}},
+    },
+    "destination": [
+        {
+            "hostname": "172.16.0.22",
+            "directory": "/home/application/testFiles/dest",
+            "protocol": {"name": "sftp", "credentials": {"username": "application"}},
+        },
+    ],
+}
+
 
 sftp_task_definition_no_permissions = {
     "type": "transfer",
@@ -335,6 +353,19 @@ def test_sftp_basic(root_dir, setup_sftp_keys):
     )
 
 
+def test_sftp_basic_non_existent_source_directory(root_dir, setup_sftp_keys):
+    # Create a transfer object
+    transfer_obj = transfer.Transfer(
+        None,
+        "sftp-basic-non-existent-directory",
+        sftp_task_non_existent_source_dir_definition,
+    )
+
+    # Run the transfer and expect a true status
+    with pytest.raises(exceptions.FilesDoNotMeetConditionsError):
+        transfer_obj.run()
+
+
 def test_sftp_basic_no_permissions(root_dir, setup_sftp_keys):
     # Create a test file
     fs.create_files(