Convert ConfigDriveHelper to being a context manager itself.

This is a little bit less awkward than the previous implementation,
and is based on a suggestion from Padraig Brady.

Change-Id: If8b09a41d21777a4bd56af16e1f7e02d61d1dd4e

4 files changed, 16 insertions, 15 deletions

diff --git a/nova/virt/configdrive.py b/nova/virt/configdrive.py
index 9255ab851..d4352c5e6 100644
--- a/nova/virt/configdrive.py
+++ b/nova/virt/configdrive.py
@@ -17,7 +17,6 @@
 
 """Config Drive v2 helper."""
 
-import contextlib
 import os
 import shutil
 import tempfile
@@ -58,17 +57,8 @@ CONF.register_opts(configdrive_opts)
 CONFIGDRIVESIZE_BYTES = 64 * 1024 * 1024
 
 
-@contextlib.contextmanager
-def config_drive_helper(instance_md=None):
-    cdb = _ConfigDriveBuilder(instance_md=instance_md)
-    try:
-        yield cdb
-    finally:
-        cdb.cleanup()
-
-
-class _ConfigDriveBuilder(object):
-    """Don't use this directly, use the fancy pants contextlib helper above!"""
+class ConfigDriveBuilder(object):
+    """Build config drives, optionally as a context manager."""
 
     def __init__(self, instance_md=None):
         self.imagefile = None
@@ -82,6 +72,17 @@ class _ConfigDriveBuilder(object):
         if instance_md is not None:
             self.add_instance_metadata(instance_md)
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exctype, excval, exctb):
+        if exctype is not None:
+            # NOTE(mikal): this means we're being cleaned up because an
+            # exception was thrown. All bets are off now, and we should not
+            # swallow the exception
+            return False
+        self.cleanup()
+
     def _add_file(self, path, data):
         filepath = os.path.join(self.tempdir, path)
         dirname = os.path.dirname(filepath)
diff --git a/nova/virt/hyperv/vmops.py b/nova/virt/hyperv/vmops.py
index 83493f7ff..1fba15506 100644
--- a/nova/virt/hyperv/vmops.py
+++ b/nova/virt/hyperv/vmops.py
@@ -192,7 +192,7 @@ class VMOps(baseops.BaseOps):
         LOG.info(_('Creating config drive at %(path)s'),
                  {'path': configdrive_path_iso}, instance=instance)
 
-        with configdrive.config_drive_helper(instance_md=inst_md) as cdb:
+        with configdrive.ConfigDriveBuilder(instance_md=inst_md) as cdb:
             try:
                 cdb.make_drive(configdrive_path_iso)
             except exception.ProcessExecutionError, e:
diff --git a/nova/virt/libvirt/driver.py b/nova/virt/libvirt/driver.py
index 83d2c96bd..42d9dd99b 100644
--- a/nova/virt/libvirt/driver.py
+++ b/nova/virt/libvirt/driver.py
@@ -1385,7 +1385,7 @@ class LibvirtDriver(driver.ComputeDriver):
 
             inst_md = instance_metadata.InstanceMetadata(instance,
                 content=files, extra_md=extra_md)
-            with configdrive.config_drive_helper(instance_md=inst_md) as cdb:
+            with configdrive.ConfigDriveBuilder(instance_md=inst_md) as cdb:
                 configdrive_path = basepath(fname='disk.config')
                 LOG.info(_('Creating config drive at %(path)s'),
                          {'path': configdrive_path}, instance=instance)
diff --git a/nova/virt/xenapi/vm_utils.py b/nova/virt/xenapi/vm_utils.py
index ac35a4f2b..6a0116098 100644
--- a/nova/virt/xenapi/vm_utils.py
+++ b/nova/virt/xenapi/vm_utils.py
@@ -894,7 +894,7 @@ def generate_configdrive(session, instance, vm_ref, userdevice,
             inst_md = instance_metadata.InstanceMetadata(instance,
                                                          content=files,
                                                          extra_md=extra_md)
-            with configdrive.config_drive_helper(instance_md=inst_md) as cdb:
+            with configdrive.ConfigDriveBuilder(instance_md=inst_md) as cdb:
                 with utils.tempdir() as tmp_path:
                     tmp_file = os.path.join(tmp_path, 'configdrive')
                     cdb.make_drive(tmp_file)