From 219a81e24720d0625c2a14ffe25a0621532458f1 Mon Sep 17 00:00:00 2001
From: LegenJCdary <jay_cob@interia.pl>
Date: Fri, 20 May 2022 12:59:28 +0200
Subject: [PATCH] Add --no-preserve option to cleanup workdir

Add test for this case
Add new testing function (check_file_startingwith_not_in_dir) in testing lib
---
 source/main.py       | 10 ++++++++++
 tests/02-checkrun.sh |  5 +++++
 tests/testing_lib.sh | 11 +++++++++++
 3 files changed, 26 insertions(+)

diff --git a/source/main.py b/source/main.py
index 761c829b..e0106f56 100644
--- a/source/main.py
+++ b/source/main.py
@@ -5,6 +5,7 @@
 import argparse
 import datetime
 import errno
+from shutil import rmtree
 import pkg_resources
 from ansible_deployer.modules.globalvars import SUBCOMMANDS
 from ansible_deployer.modules.configs.config import Config
@@ -49,6 +50,8 @@ def parse_options(argv):
                         help='Setup repo outside of workdir in requested path. This option applies'
                         ' only to infrastructures with allow_user_checkout enabled in infra'
                         ' config!')
+    parser.add_argument("--no-preserve", default=False, action="store_true", help='Remove workdir '
+                        ' after succesful execution.')
 
     arguments = parser.parse_args(argv)
 
@@ -87,6 +90,7 @@ def parse_options(argv):
         options["conf_dir"] = os.path.abspath(arguments.conf_dir[0])
     else:
         options["conf_dir"] = None
+    options["no_preserve"] = arguments.no_preserve
 
     return options
 
@@ -143,6 +147,12 @@ def main():
             lock.lock_inventory(lockdir, lockpath)
             runner.run_playitem(config, options, inv_file, lockpath)
             lock.unlock_inventory(lockpath)
+            if options["no_preserve"]:
+                try:
+                    rmtree(workdir)
+                    logger.logger.info("Working directory succesfully removed.")
+                except Exception as exc:
+                    logger.logger.critical("Failed to remove working directory due to: %s.", exc)
         elif options["subcommand"] == "lock":
             lock.lock_inventory(lockdir, lockpath)
         elif options["subcommand"] == "unlock":
diff --git a/tests/02-checkrun.sh b/tests/02-checkrun.sh
index e4422ac4..57e103f8 100755
--- a/tests/02-checkrun.sh
+++ b/tests/02-checkrun.sh
@@ -79,6 +79,11 @@ check_message_in_output 'ansible-deployer show task' 'Available tasks:'
 check_run_ok "ansible-deployer run -t task_with_multi_groups -s testing -i testInfra"
 check_message_in_output "ansible-deployer run -t task_with_multi_groups_fail -s testing -i testInfra" "\[CRITICAL\]: Task forbidden"
 
+# Check --no-preserve option
+check_run_ok "ansible-deployer run -t task_exec_bin_true -s prod -i testInfra --no-preserve"
+search_path=$(find_latest_sequence)
+check_file_startingwith_not_in_dir "$search_path" "ansible-deploy_execution_"
+
 #Try execution of task without permissions
 if [ $UID -ne 0 ]
 then
diff --git a/tests/testing_lib.sh b/tests/testing_lib.sh
index 2c8a7fe7..edd4d768 100644
--- a/tests/testing_lib.sh
+++ b/tests/testing_lib.sh
@@ -72,3 +72,14 @@ check_file_startingwith_in_dir() {
 		exit 1
 	fi
 }
+
+check_file_startingwith_not_in_dir() {
+        eval "find $1 -name "${2}*" 1>/dev/null"
+        if [ $? -eq 0 ]
+        then
+                echo "FAILED: $1 has file starting with pattern $2"
+        else
+                echo "OK: $1 does not contain file starting with pattern $2"
+                exit 1
+        fi
+}