Merge remote-tracking branch 'origin/baserock/richardmaw/S8544/deploy-cleanup-v3'

Reviewed-by: Lars Wirzenius
Reviewed-by: Daniel Silverstone

3 files changed, 106 insertions, 55 deletions

diff --git a/morphlib/plugins/deploy_plugin.py b/morphlib/plugins/deploy_plugin.py
index 4e588eaa..8530cb57 100644
--- a/morphlib/plugins/deploy_plugin.py
+++ b/morphlib/plugins/deploy_plugin.py
@@ -225,6 +225,9 @@ class DeployPlugin(cliapp.Plugin):
         location = args[2]
         env_vars = args[3:]
 
+        # Set up environment for running extensions.
+        env = morphlib.util.parse_environment_pairs(os.environ, env_vars)
+
         # Deduce workspace and system branch and branch root repository.
         workspace = self.other.deduce_workspace()
         branch, branch_dir = self.other.deduce_system_branch()
@@ -280,67 +283,65 @@ class DeployPlugin(cliapp.Plugin):
         if push:
             self.other.delete_remote_build_refs(build_repos)
             
-        # Unpack the artifact (tarball) to a temporary directory.
-        self.app.status(msg='Unpacking system for configuration')
 
-        system_tree = tempfile.mkdtemp(
+        # Create a tempdir for this deployment to work in
+        deploy_tempdir = tempfile.mkdtemp(
             dir=os.path.join(self.app.settings['tempdir'], 'deployments'))
-
-        if build_command.lac.has(artifact):
-            f = build_command.lac.get(artifact)
-        elif build_command.rac.has(artifact):
-            f = build_command.rac.get(artifact)
-        else:
-            raise cliapp.AppException('Deployment failed as system is not yet'
-                                      ' built.\nPlease ensure system is built'
-                                      ' before deployment.')
-        tf = tarfile.open(fileobj=f)
-        tf.extractall(path=system_tree)
-        
-        self.app.status(
-            msg='System unpacked at %(system_tree)s',
-            system_tree=system_tree)
-
-        # Set up environment for running extensions.
-        env = dict(os.environ)
-        for spec in env_vars:
-            name, value = spec.split('=', 1)
-            if name in env:
-                raise morphlib.Error(
-                    '%s is already set in the enviroment' % name)
-            env[name] = value
-
-        if 'TMPDIR' not in env:
-            # morphlib.app already took care of ensuring the tempdir setting
-            # is good, so use it if we don't have one already set.
-            env['TMPDIR'] = os.path.join(self.app.settings['tempdir'],
-                                         'deployments')
-
-        # Run configuration extensions.
-        self.app.status(msg='Configure system')
-        names = artifact.source.morphology['configuration-extensions']
-        for name in names:
+        try:
+            # Create a tempdir to extract the rootfs in
+            system_tree = tempfile.mkdtemp(dir=deploy_tempdir)
+
+            # Extensions get a private tempdir so we can more easily clean
+            # up any files an extension left behind
+            deploy_private_tempdir = tempfile.mkdtemp(dir=deploy_tempdir)
+            env['TMPDIR'] = deploy_private_tempdir
+
+            # Unpack the artifact (tarball) to a temporary directory.
+            self.app.status(msg='Unpacking system for configuration')
+
+            if build_command.lac.has(artifact):
+                f = build_command.lac.get(artifact)
+            elif build_command.rac.has(artifact):
+                f = build_command.rac.get(artifact)
+            else:
+                raise cliapp.AppException('Deployment failed as system is'
+                                          ' not yet built.\nPlease ensure'
+                                          ' the system is built before'
+                                          ' deployment.')
+            tf = tarfile.open(fileobj=f)
+            tf.extractall(path=system_tree)
+
+            self.app.status(
+                msg='System unpacked at %(system_tree)s',
+                system_tree=system_tree)
+
+
+            # Run configuration extensions.
+            self.app.status(msg='Configure system')
+            names = artifact.source.morphology['configuration-extensions']
+            for name in names:
+                self._run_extension(
+                    root_repo_dir,
+                    build_ref,
+                    name,
+                    '.configure',
+                    [system_tree],
+                    env)
+
+            # Run write extension.
+            self.app.status(msg='Writing to device')
             self._run_extension(
                 root_repo_dir,
                 build_ref,
-                name,
-                '.configure',
-                [system_tree],
+                deployment_type,
+                '.write',
+                [system_tree, location],
                 env)
-        
-        # Run write extension.
-        self.app.status(msg='Writing to device')
-        self._run_extension(
-            root_repo_dir,
-            build_ref,
-            deployment_type,
-            '.write',
-            [system_tree, location],
-            env)
-        
-        # Cleanup.
-        self.app.status(msg='Cleaning up')
-        shutil.rmtree(system_tree)
+
+        finally:
+            # Cleanup.
+            self.app.status(msg='Cleaning up')
+            shutil.rmtree(deploy_tempdir)
 
         self.app.status(msg='Finished deployment')
 
diff --git a/morphlib/util.py b/morphlib/util.py
index ead0bafe..b83211e3 100644
--- a/morphlib/util.py
+++ b/morphlib/util.py
@@ -327,3 +327,35 @@ def find_leaf(dirname, subdir_name):
             return None
         dirname = subdirs[0]
 
+
+class EnvironmentAlreadySetError(morphlib.Error):
+
+    def __init__(self, conflicts):
+        self.conflicts = conflicts
+        morphlib.Error.__init__(
+            self, 'Keys %r are already set in the environment' % conflicts)
+
+
+def parse_environment_pairs(env, pairs):
+    '''Add key=value pairs to the environment dict.
+
+    Given a dict and a list of strings of the form key=value,
+    set dict[key] = value, unless key is already set in the
+    environment, at which point raise an exception.
+
+    This does not modify the passed in dict.
+
+    Returns the extended dict.
+
+    '''
+
+    extra_env = dict(p.split('=', 1) for p in pairs)
+    conflicting = [k for k in extra_env if k in env]
+    if conflicting:
+        raise EnvironmentAlreadySetError(conflicting)
+
+    # Return a dict that is the union of the two
+    # This is not the most performant, since it creates
+    # 3 unnecessary lists, but I felt this was the most
+    # easy to read. Using itertools.chain may be more efficicent
+    return dict(env.items() + extra_env.items())
diff --git a/morphlib/util_tests.py b/morphlib/util_tests.py
index eaff0821..ca9fe5ae 100644
--- a/morphlib/util_tests.py
+++ b/morphlib/util_tests.py
@@ -87,3 +87,21 @@ class FindParentOfTests(unittest.TestCase):
     def test_find_leaf_returns_none_if_not_found(self):
         self.assertEqual(morphlib.util.find_leaf(self.a, '.magic'), None)
 
+
+class ParseEnvironmentPairsTests(unittest.TestCase):
+
+    def test_parse_environment_pairs_adds_key(self):
+        ret = morphlib.util.parse_environment_pairs({}, ["foo=bar"])
+        self.assertEqual(ret.get("foo"), "bar")
+
+    def test_parse_environment_does_not_alter_passed_dict(self):
+        d = {}
+        morphlib.util.parse_environment_pairs(d, ["foo=bar"])
+        self.assertTrue("foo" not in d)
+
+    def test_parse_environment_raises_on_duplicates(self):
+        self.assertRaises(
+            morphlib.util.EnvironmentAlreadySetError,
+            morphlib.util.parse_environment_pairs,
+            {"foo": "bar"},
+            ["foo=bar"])