Change path removal suitable for integration and normal run.

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)

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,