Add utils.tempdir() context manager for easy temp dirs

Fixes bug 883323 (and others)

Users of tempfile.mkdtemp() need to make sure the directory is cleaned
up when it's done being used. Unfortunately, not all of the code does
so at all, or safely (by using a try/finally block).

Change-Id: I270109d83efec4f8b3dd954021493f4d96c6ab79

nova/auth/manager.py

@@ -24,9 +24,7 @@ Nova authentication management
 """
 
 import os
-import shutil
 import string  # pylint: disable=W0402
-import tempfile
 import uuid
 import zipfile
 
@@ -767,7 +765,7 @@ class AuthManager(object):
         pid = Project.safe_id(project)
         private_key, signed_cert = crypto.generate_x509_cert(user.id, pid)
 
-        tmpdir = tempfile.mkdtemp()
+        with utils.tempdir() as tmpdir:
             zf = os.path.join(tmpdir, "temp.zip")
             zippy = zipfile.ZipFile(zf, 'w')
             if use_dmz and FLAGS.region_list:
@@ -805,7 +803,6 @@ class AuthManager(object):
             with open(zf, 'rb') as f:
                 read_buffer = f.read()
 
-        shutil.rmtree(tmpdir)
         return read_buffer
 
     def get_environment_rc(self, user, project=None, use_dmz=True):

nova/compat/flagfile.py

@@ -175,6 +175,8 @@ def handle_flagfiles_managed(args):
              # Do stuff
         # Any temporary fils have been removed
     '''
+    # NOTE(johannes): Would be nice to use utils.tempdir(), but it
+    # causes an import loop
     tempdir = tempfile.mkdtemp(prefix='nova-conf-')
     try:
         yield handle_flagfiles(args, tempdir=tempdir)

nova/tests/test_imagecache.py

@@ -17,12 +17,11 @@
 #    under the License.
 
 
+import contextlib
 import cStringIO
 import hashlib
 import logging
 import os
-import shutil
-import tempfile
 import time
 
 from nova import test
@@ -58,9 +57,8 @@ class ImageCacheManagerTestCase(test.TestCase):
         self.assertEquals(csum, None)
 
     def test_read_stored_checksum(self):
-        try:
+        with utils.tempdir() as tmpdir:
-            dirname = tempfile.mkdtemp()
+            fname = os.path.join(tmpdir, 'aaa')
-            fname = os.path.join(dirname, 'aaa')
 
             csum_input = 'fdghkfhkgjjksfdgjksjkghsdf'
             f = open('%s.sha1' % fname, 'w')
@@ -71,9 +69,6 @@ class ImageCacheManagerTestCase(test.TestCase):
 
             self.assertEquals(csum_input, csum_output)
 
-        finally:
-            shutil.rmtree(dirname)
-
     def test_list_base_images(self):
         listing = ['00000001',
                    'ephemeral_0_20_None',
@@ -281,13 +276,17 @@ class ImageCacheManagerTestCase(test.TestCase):
                                 (base_file2, True, False),
                                 (base_file3, False, True)])
 
+    @contextlib.contextmanager
     def _intercept_log_messages(self):
+        try:
             mylog = log.getLogger()
             stream = cStringIO.StringIO()
             handler = logging.StreamHandler(stream)
             handler.setFormatter(log.LegacyNovaFormatter())
             mylog.logger.addHandler(handler)
-        return mylog, handler, stream
+            yield stream
+        finally:
+            mylog.logger.removeHandler(handler)
 
     def test_verify_checksum(self):
         testdata = ('OpenStack Software delivers a massively scalable cloud '
@@ -295,11 +294,10 @@ class ImageCacheManagerTestCase(test.TestCase):
         img = {'container_format': 'ami', 'id': '42'}
 
         self.flags(checksum_base_images=True)
-        mylog, handler, stream = self._intercept_log_messages()
 
-        try:
+        with self._intercept_log_messages() as stream:
-            dirname = tempfile.mkdtemp()
+            with utils.tempdir() as tmpdir:
-            fname = os.path.join(dirname, 'aaa')
+                fname = os.path.join(tmpdir, 'aaa')
 
                 f = open(fname, 'w')
                 f.write(testdata)
@@ -324,8 +322,8 @@ class ImageCacheManagerTestCase(test.TestCase):
                 image_cache_manager = imagecache.ImageCacheManager()
                 res = image_cache_manager._verify_checksum(img, fname)
                 self.assertFalse(res)
-            self.assertNotEqual(stream.getvalue().find('image verification '
+                log = stream.getvalue()
-                                                       'failed'), -1)
+                self.assertNotEqual(log.find('image verification failed'), -1)
 
                 # Checksum file missing
                 os.remove('%s.sha1' % fname)
@@ -337,15 +335,12 @@ class ImageCacheManagerTestCase(test.TestCase):
                 # side effect of creating the checksum
                 self.assertTrue(os.path.exists('%s.sha1' % fname))
 
-        finally:
+    @contextlib.contextmanager
-            shutil.rmtree(dirname)
+    def _make_base_file(self, checksum=True):
-            mylog.logger.removeHandler(handler)
-
-    def _make_base_file(checksum=True):
         """Make a base file for testing."""
 
-        dirname = tempfile.mkdtemp()
+        with utils.tempdir() as tmpdir:
-        fname = os.path.join(dirname, 'aaa')
+            fname = os.path.join(tmpdir, 'aaa')
 
             base_file = open(fname, 'w')
             base_file.write('data')
@@ -358,11 +353,10 @@ class ImageCacheManagerTestCase(test.TestCase):
                 checksum_file.close()
 
             base_file.close()
-        return dirname, fname
+            yield fname
 
     def test_remove_base_file(self):
-        dirname, fname = self._make_base_file()
+        with self._make_base_file() as fname:
-        try:
             image_cache_manager = imagecache.ImageCacheManager()
             image_cache_manager._remove_base_file(fname)
 
@@ -377,12 +371,8 @@ class ImageCacheManagerTestCase(test.TestCase):
             self.assertFalse(os.path.exists(fname))
             self.assertFalse(os.path.exists('%s.sha1' % fname))
 
-        finally:
-            shutil.rmtree(dirname)
-
     def test_remove_base_file_original(self):
-        dirname, fname = self._make_base_file()
+        with self._make_base_file() as fname:
-        try:
             image_cache_manager = imagecache.ImageCacheManager()
             image_cache_manager.originals = [fname]
             image_cache_manager._remove_base_file(fname)
@@ -405,27 +395,19 @@ class ImageCacheManagerTestCase(test.TestCase):
             self.assertFalse(os.path.exists(fname))
             self.assertFalse(os.path.exists('%s.sha1' % fname))
 
-        finally:
-            shutil.rmtree(dirname)
-
     def test_remove_base_file_dne(self):
         # This test is solely to execute the "does not exist" code path. We
         # don't expect the method being tested to do anything in this case.
-        dirname = tempfile.mkdtemp()
+        with utils.tempdir() as tmpdir:
-        try:
+            fname = os.path.join(tmpdir, 'aaa')
-            fname = os.path.join(dirname, 'aaa')
             image_cache_manager = imagecache.ImageCacheManager()
             image_cache_manager._remove_base_file(fname)
 
-        finally:
-            shutil.rmtree(dirname)
-
     def test_remove_base_file_oserror(self):
-        dirname = tempfile.mkdtemp()
+        with self._intercept_log_messages() as stream:
-        fname = os.path.join(dirname, 'aaa')
+            with utils.tempdir() as tmpdir:
-        mylog, handler, stream = self._intercept_log_messages()
+                fname = os.path.join(tmpdir, 'aaa')
 
-        try:
                 os.mkdir(fname)
                 os.utime(fname, (-1, time.time() - 3601))
 
@@ -437,19 +419,14 @@ class ImageCacheManagerTestCase(test.TestCase):
                 self.assertNotEqual(stream.getvalue().find('Failed to remove'),
                                     -1)
 
-        finally:
-            shutil.rmtree(dirname)
-            mylog.logger.removeHandler(handler)
-
     def test_handle_base_image_unused(self):
         img = {'container_format': 'ami',
                'id': '123',
                'uuid': '1234-4567-2378'}
 
-        dirname, fname = self._make_base_file()
+        with self._make_base_file() as fname:
             os.utime(fname, (-1, time.time() - 3601))
 
-        try:
             image_cache_manager = imagecache.ImageCacheManager()
             image_cache_manager.unexplained_images = [fname]
             image_cache_manager._handle_base_image(img, fname)
@@ -459,18 +436,14 @@ class ImageCacheManagerTestCase(test.TestCase):
                               [fname])
             self.assertEquals(image_cache_manager.corrupt_base_files, [])
 
-        finally:
-            shutil.rmtree(dirname)
-
     def test_handle_base_image_used(self):
         img = {'container_format': 'ami',
                'id': '123',
                'uuid': '1234-4567-2378'}
 
-        dirname, fname = self._make_base_file()
+        with self._make_base_file() as fname:
             os.utime(fname, (-1, time.time() - 3601))
 
-        try:
             image_cache_manager = imagecache.ImageCacheManager()
             image_cache_manager.unexplained_images = [fname]
             image_cache_manager.used_images = {'123': (1, 0, ['banana-42'])}
@@ -480,18 +453,14 @@ class ImageCacheManagerTestCase(test.TestCase):
             self.assertEquals(image_cache_manager.removable_base_files, [])
             self.assertEquals(image_cache_manager.corrupt_base_files, [])
 
-        finally:
-            shutil.rmtree(dirname)
-
     def test_handle_base_image_used_remotely(self):
         img = {'container_format': 'ami',
                'id': '123',
                'uuid': '1234-4567-2378'}
 
-        dirname, fname = self._make_base_file()
+        with self._make_base_file() as fname:
             os.utime(fname, (-1, time.time() - 3601))
 
-        try:
             image_cache_manager = imagecache.ImageCacheManager()
             image_cache_manager.used_images = {'123': (0, 1, ['banana-42'])}
             image_cache_manager._handle_base_image(img, None)
@@ -500,9 +469,6 @@ class ImageCacheManagerTestCase(test.TestCase):
             self.assertEquals(image_cache_manager.removable_base_files, [])
             self.assertEquals(image_cache_manager.corrupt_base_files, [])
 
-        finally:
-            shutil.rmtree(dirname)
-
     def test_handle_base_image_absent(self):
         """Ensure we warn for use of a missing base image."""
 
@@ -510,9 +476,7 @@ class ImageCacheManagerTestCase(test.TestCase):
                'id': '123',
                'uuid': '1234-4567-2378'}
 
-        mylog, handler, stream = self._intercept_log_messages()
+        with self._intercept_log_messages() as stream:
-
-        try:
             image_cache_manager = imagecache.ImageCacheManager()
             image_cache_manager.used_images = {'123': (1, 0, ['banana-42'])}
             image_cache_manager._handle_base_image(img, None)
@@ -523,18 +487,14 @@ class ImageCacheManagerTestCase(test.TestCase):
             self.assertNotEqual(stream.getvalue().find('an absent base file'),
                                 -1)
 
-        finally:
-            mylog.logger.removeHandler(handler)
-
     def test_handle_base_image_used_missing(self):
         img = {'container_format': 'ami',
                'id': '123',
                'uuid': '1234-4567-2378'}
 
-        dirname = tempfile.mkdtemp()
+        with utils.tempdir() as tmpdir:
-        fname = os.path.join(dirname, 'aaa')
+            fname = os.path.join(tmpdir, 'aaa')
 
-        try:
             image_cache_manager = imagecache.ImageCacheManager()
             image_cache_manager.unexplained_images = [fname]
             image_cache_manager.used_images = {'123': (1, 0, ['banana-42'])}
@@ -544,17 +504,12 @@ class ImageCacheManagerTestCase(test.TestCase):
             self.assertEquals(image_cache_manager.removable_base_files, [])
             self.assertEquals(image_cache_manager.corrupt_base_files, [])
 
-        finally:
-            shutil.rmtree(dirname)
-
     def test_handle_base_image_checksum_fails(self):
         img = {'container_format': 'ami',
                'id': '123',
                'uuid': '1234-4567-2378'}
 
-        dirname, fname = self._make_base_file()
+        with self._make_base_file() as fname:
-
-        try:
             f = open(fname, 'w')
             f.write('banana')
             f.close()
@@ -569,9 +524,6 @@ class ImageCacheManagerTestCase(test.TestCase):
             self.assertEquals(image_cache_manager.corrupt_base_files,
                               [fname])
 
-        finally:
-            shutil.rmtree(dirname)
-
     def test_verify_base_images(self):
         self.flags(instances_path='/instance_path')
         self.flags(remove_unused_base_images=True)

nova/tests/test_libvirt.py

@@ -1050,7 +1050,7 @@ class LibvirtConnTestCase(test.TestCase):
     def test_pre_block_migration_works_correctly(self):
         """Confirms pre_block_migration works correctly."""
         # Replace instances_path since this testcase creates tmpfile
-        tmpdir = tempfile.mkdtemp()
+        with utils.tempdir() as tmpdir:
             self.flags(instances_path=tmpdir)
 
             # Test data
@@ -1069,7 +1069,6 @@ class LibvirtConnTestCase(test.TestCase):
             self.assertTrue(os.path.exists('%s/%s/' %
                                            (tmpdir, instance_ref.name)))
 
-        shutil.rmtree(tmpdir)
         db.instance_destroy(self.context, instance_ref['id'])
 
     @test.skip_if(missing_libvirt(), "Test requires libvirt")
@@ -1926,13 +1925,10 @@ disk size: 4.4M''', ''))
         libvirt_utils.mkfs('swap', '/my/swap/block/dev')
 
     def test_ensure_tree(self):
-        tmpdir = tempfile.mkdtemp()
+        with utils.tempdir() as tmpdir:
-        try:
             testdir = '%s/foo/bar/baz' % (tmpdir,)
             libvirt_utils.ensure_tree(testdir)
             self.assertTrue(os.path.isdir(testdir))
-        finally:
-            shutil.rmtree(tmpdir)
 
     def test_write_to_file(self):
         dst_fd, dst_path = tempfile.mkstemp()