From ac04edc233c61b8b730ebc43fe3c63e78d8d4cc4 Mon Sep 17 00:00:00 2001
From: Darko Poljak <foss@ungleich.com>
Date: Thu, 14 Sep 2017 09:53:16 +0200
Subject: [PATCH] Change path removal suitable for integration and normal run.

---
 cdist/config.py      | 19 ++++++++++++++++++-
 cdist/integration.py |  2 ++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/cdist/config.py b/cdist/config.py
index 11da04da..f5c3121d 100644
--- a/cdist/config.py
+++ b/cdist/config.py
@@ -44,6 +44,22 @@ from cdist.util.remoteutil import inspect_ssh_mux_opts
 class Config(object):
     """Cdist main class to hold arbitrary data"""
 
+    # list of paths (files and/or directories) that will be removed on finish
+    _paths_for_removal = []
+
+    @classmethod
+    def _register_path_for_removal(cls, path):
+        cls._paths_for_removal.append(path)
+
+    @classmethod
+    def _remove_paths(cls):
+        while cls._paths_for_removal:
+            path = cls._paths_for_removal.pop()
+            if os.path.isfile(path):
+                os.remove(path)
+            else:
+                shutil.rmtree(path)
+
     def __init__(self, local, remote, dry_run=False, jobs=None,
                  cleanup_cmds=None, remove_remote_files_dirs=False):
 
@@ -266,8 +282,8 @@ class Config(object):
     @classmethod
     def _resolve_ssh_control_path(cls):
         base_path = tempfile.mkdtemp()
+        cls._register_path_for_removal(base_path)
         control_path = os.path.join(base_path, "s")
-        atexit.register(lambda: shutil.rmtree(base_path))
         return control_path
 
     @classmethod
@@ -345,6 +361,7 @@ class Config(object):
                     cleanup_cmds=cleanup_cmds,
                     remove_remote_files_dirs=remove_remote_files_dirs)
             c.run()
+            cls._remove_paths()
 
         except cdist.Error as e:
             log.error(e)
diff --git a/cdist/integration.py b/cdist/integration.py
index e2f9be6e..ee742cc5 100644
--- a/cdist/integration.py
+++ b/cdist/integration.py
@@ -36,6 +36,7 @@ import os
 import os.path
 import collections
 import uuid
+import shutil
 
 
 def find_cdist_exec_in_path():
@@ -123,6 +124,7 @@ def _process_hosts_simple(action, host, manifest, verbose,
         theclass.onehost(target_host, None, host_base_path, hostdir, args,
                          parallel=False, configuration=configuration,
                          remove_remote_files_dirs=True)
+    shutil.rmtree(base_root_path)
 
 
 def configure_hosts_simple(host, manifest,