Refactor MetadataFileContext to DRY things up

This appears to have been copied over to pulp_rpm at some point
and subsequently modified to have yum-specific behavior. Behaviors
that were not yum-specific and the related tests have been pushed
back to platform.

server/pulp/plugins/util/metadata_writer.py

@@ -1,47 +1,51 @@
-from gettext import gettext as _
 import glob
 import gzip
+import hashlib
 import logging
 import os
 import shutil
 import traceback
-
-
+from gettext import gettext as _
+from xml.dom import minidom
 from xml.sax.saxutils import XMLGenerator
 
 from pulp.common import error_codes
-from pulp.server.exceptions import PulpCodedValidationException, PulpCodedException
-from verification import CHECKSUM_FUNCTIONS
+from pulp.server.exceptions import PulpCodedException, PulpCodedValidationException
 
 _LOG = logging.getLogger(__name__)
 BUFFER_SIZE = 1024
+HASHLIB_ALGORITHMS = ('md5', 'sha1', 'sha224', 'sha256', 'sha384', 'sha512')
 
 
 class MetadataFileContext(object):
     """
     Context manager class for metadata file generation.
     """
 
-    def __init__(self, metadata_file_path, checksum_type=None):
+    def __init__(self, metadata_file_path, checksum_type=None, no_checksum_filenames=None):
         """
         :param metadata_file_path: full path to metadata file to be generated
         :type  metadata_file_path: str
         :param checksum_type: checksum type to be used to generate and prepend checksum
-                              to the file names of files. If checksum_type is None,
+                              to the file names of metadata files. If checksum_type is None,
                               no checksum is added to the filename
         :type checksum_type: str or None
+        :param no_checksum_filenames: file names that *will not* have a checksum prepended to them
+        :type no_checksum_filenames: list or None
         """
 
         self.metadata_file_path = metadata_file_path
         self.metadata_file_handle = None
         self.checksum_type = checksum_type
         self.checksum = None
         if self.checksum_type is not None:
-            checksum_function = CHECKSUM_FUNCTIONS.get(checksum_type)
-            if not checksum_function:
-                raise PulpCodedValidationException(
-                    [PulpCodedException(error_codes.PLP1005, checksum_type=checksum_type)])
-            self.checksum_constructor = checksum_function
+            try:
+                self.checksum_constructor = getattr(hashlib, checksum_type)
+            except AttributeError:
+                raise PulpCodedValidationException([PulpCodedException(error_codes.PLP1005,
+                                                    checksum_type=checksum_type)])
+        if no_checksum_filenames is None:
+            self.no_checksum_filenames = []
 
     def __enter__(self):
 
@@ -52,7 +56,6 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_val, exc_tb):
 
         if None not in (exc_type, exc_val, exc_tb):
-
             err_msg = '\n'.join(traceback.format_exception(exc_type, exc_val, exc_tb))
             log_msg = _('Exception occurred while writing [%(m)s]\n%(e)s')
             # any errors here should have already been caught and logged
@@ -64,25 +67,27 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
     def initialize(self):
         """
-        Create the new metadata file and write the header.
+        Create a new metadata file and write the XML header and opening root
+        level tag into it.
         """
         if self.metadata_file_handle is not None:
             # initialize has already, at least partially, been run
             return
 
         self._open_metadata_file_handle()
-        self._write_file_header()
+        self._write_xml_header()
+        self._write_root_tag_open()
 
     def finalize(self):
         """
-        Write the footer into the metadata file and close it.
+        Write the closing root level tag into the metadata file and close it.
         """
         if self._is_closed(self.metadata_file_handle):
             # finalize has already been run or initialize has not been run
             return
 
         try:
-            self._write_file_footer()
+            self._write_root_tag_close()
 
         except Exception, e:
             _LOG.exception(e)
@@ -93,9 +98,9 @@ def finalize(self):
         except Exception, e:
             _LOG.exception(e)
 
-        # Add calculated checksum to the filename
+        # Add calculated checksum to filenames not in no_checksum_filenames
         file_name = os.path.basename(self.metadata_file_path)
-        if self.checksum_type is not None:
+        if self.checksum_type is not None and file_name not in self.no_checksum_filenames:
             with open(self.metadata_file_path, 'rb') as file_handle:
                 content = file_handle.read()
                 checksum = self.checksum_constructor(content).hexdigest()
@@ -107,9 +112,6 @@ def finalize(self):
             os.rename(self.metadata_file_path, new_file_path)
             self.metadata_file_path = new_file_path
 
-        # Set the metadata_file_handle to None so we don't double call finalize
-        self.metadata_file_handle = None
-
     def _open_metadata_file_handle(self):
         """
         Open the metadata file handle, creating any missing parent directories.
@@ -148,17 +150,26 @@ def _open_metadata_file_handle(self):
         else:
             self.metadata_file_handle = open(self.metadata_file_path, 'w')
 
-    def _write_file_header(self):
+    def _write_xml_header(self):
         """
-        Write any headers for the metadata file
+        Write the initial <?xml?> header tag into the file handle.
         """
-        pass
+        assert self.metadata_file_handle is not None
+        _LOG.debug('Writing XML header into metadata file: %s' % self.metadata_file_path)
+        xml_header = minidom.Document().toxml(encoding='UTF-8') + os.linesep
+        self.metadata_file_handle.write(xml_header)
 
-    def _write_file_footer(self):
+    def _write_root_tag_open(self):
+        """
+        Write the opening tag for the root element of a given metadata XML file.
+        """
+        raise NotImplementedError()
+
+    def _write_root_tag_close(self):
         """
-        Write any file footers for the metadata file.
+        Write the closing tag for the root element of a give metadata XML file.
         """
-        pass
+        raise NotImplementedError()
 
     def _close_metadata_file_handle(self):
         """

server/test/unit/plugins/util/test_metadata_writer.py

@@ -26,7 +26,11 @@ def setUp(self):
         self.metadata_file_dir = tempfile.mkdtemp()
 
     def tearDown(self):
-        shutil.rmtree(self.metadata_file_dir)
+        # some tests change dir permissions, so recursively
+        # fix dir permissions before cleaning up the tmp dir
+        for dirname, __, __ in os.walk(self.metadata_file_dir):
+            os.chmod(dirname, 0700)
+        os.chmod(self.metadata_file_dir, 0700)
 
     def test_metadata_instantiation(self):
         try:
@@ -67,7 +71,6 @@ def test_open_handle_bad_parent_permissions(self):
 
         os.makedirs(parent_path, mode=0000)
         self.assertRaises(RuntimeError, context._open_metadata_file_handle)
-        os.chmod(parent_path, 0777)
 
     def test_open_handle_file_exists(self):
 
@@ -77,7 +80,12 @@ def test_open_handle_file_exists(self):
         with open(path, 'w') as h:
             h.flush()
 
-        context._open_metadata_file_handle()
+        try:
+            context._open_metadata_file_handle()
+        except Exception, e:
+            self.fail(e.message)
+        else:
+            context._close_metadata_file_handle()
 
     def test_open_handle_bad_file_permissions(self):
 
@@ -99,7 +107,7 @@ def test_open_handle_gzip(self):
 
         self.assertTrue(os.path.exists(path))
 
-        context._write_file_header()
+        context._write_xml_header()
         context._close_metadata_file_handle()
 
         try:
@@ -110,6 +118,23 @@ def test_open_handle_gzip(self):
 
         h.close()
 
+    def test_write_xml_header(self):
+
+        path = os.path.join(self.metadata_file_dir, 'header.xml')
+        context = MetadataFileContext(path)
+
+        context._open_metadata_file_handle()
+        context._write_xml_header()
+        context._close_metadata_file_handle()
+
+        self.assertTrue(os.path.exists(path))
+
+        with open(path) as h:
+            content = h.read()
+
+        expected_content = '<?xml version="1.0" encoding="UTF-8"?>\n'
+        self.assertEqual(content, expected_content)
+
     def test_init_invalid_checksum(self):
         path = os.path.join(self.metadata_file_dir, 'foo', 'header.xml')
         assert_validation_exception(MetadataFileContext, [PLP1005], path, checksum_type='invalid')
@@ -119,11 +144,11 @@ def test_initialize(self):
         path = os.path.join(self.metadata_file_dir, 'foo', 'header.xml')
         context = MetadataFileContext(path)
 
-        context._write_file_header = Mock()
+        context._write_root_tag_open = Mock()
 
         context.initialize()
 
-        context._write_file_header.assert_called_once_with()
+        context._write_root_tag_open.assert_called_once_with()
         self.assertTrue(os.path.exists(path))
 
         with open(path) as h:
@@ -134,22 +159,23 @@ def test_initialize(self):
     def test_initialize_double_call(self):
         path = os.path.join(self.metadata_file_dir, 'test.xml')
         context = MetadataFileContext(path)
+        context._write_root_tag_open = Mock()
         context.initialize()
-        context._write_file_header = Mock()
         context.initialize()
-        self.assertEquals(0, context._write_file_header.call_count)
+        # despite initialize being called twice, the root tag should only be written once
+        context._write_root_tag_open.assert_called_once_with()
         context.finalize()
 
     def test_is_closed_gzip_file(self):
-        path = os.path.join(os.path.dirname(__file__), '../../../data/foo.tar.gz')
+        path = os.path.join(DATA_DIR, 'foo.tar.gz')
 
         file_object = gzip.open(path)
         file_object.close()
 
         self.assertTrue(MetadataFileContext._is_closed(file_object))
 
     def test_is_open_gzip_file(self):
-        path = os.path.join(os.path.dirname(__file__), '../../../data/foo.tar.gz')
+        path = os.path.join(DATA_DIR, 'foo.tar.gz')
 
         file_object = gzip.open(path)
 
@@ -158,7 +184,7 @@ def test_is_open_gzip_file(self):
         file_object.close()
 
     def test_is_closed_file(self):
-        path = os.path.join(os.path.dirname(__file__), '../../../data/foo.tar.gz')
+        path = os.path.join(DATA_DIR, 'foo.tar.gz')
 
         # opening as a regular file, not with gzip
         file_object = open(path)
@@ -167,7 +193,7 @@ def test_is_closed_file(self):
         self.assertTrue(MetadataFileContext._is_closed(file_object))
 
     def test_is_open_file(self):
-        path = os.path.join(os.path.dirname(__file__), '../../../data/foo.tar.gz')
+        path = os.path.join(DATA_DIR, 'foo.tar.gz')
 
         # opening as a regular file, not with gzip
         file_object = open(path)
@@ -185,7 +211,7 @@ def test_finalize_closed_gzip_file(self):
         # this test makes sure that we can properly detect the closed state of
         # a gzip file, because on python 2.6 we have to take special measures
         # to do so.
-        path = os.path.join(os.path.dirname(__file__), '../../../data/foo.tar.gz')
+        path = os.path.join(DATA_DIR, 'foo.tar.gz')
 
         context = MetadataFileContext('/a/b/c')
         context.metadata_file_handle = gzip.open(path)
@@ -199,13 +225,12 @@ def test_finalize_checksum_type_none(self):
         path = os.path.join(self.metadata_file_dir, 'test.xml')
         context = MetadataFileContext(path)
 
-        context.initialize()
-        context._write_file_footer = Mock()
+        context._open_metadata_file_handle()
+        context._write_xml_header()
+        context._close_metadata_file_handle()
         context.finalize()
-        context._write_file_footer.assert_called_once_with()
 
         self.assertEqual(context.metadata_file_path, path)
-        self.assertEqual(context.metadata_file_handle, None)
 
     def test_finalize_error_on_write_footer(self):
         # Ensure that if the write_file_footer throws an exception we eat it so that
@@ -215,22 +240,24 @@ def test_finalize_error_on_write_footer(self):
         path = os.path.join(self.metadata_file_dir, 'test.xml')
         context = MetadataFileContext(path)
 
+        context._write_root_tag_open = Mock()
+        context._write_root_tag_close = Mock(side_effect=Exception())
         context.initialize()
-        context._write_file_footer = Mock(side_effect=Exception())
         context.finalize()
-        context._write_file_footer.assert_called_once_with()
+        context._write_root_tag_close.assert_called_once_with()
 
         self.assertEqual(context.metadata_file_path, path)
 
     def test_finalize_double_call(self):
         path = os.path.join(self.metadata_file_dir, 'test.xml')
         context = MetadataFileContext(path)
 
+        context._write_root_tag_open = Mock()
+        context._write_root_tag_close = Mock()
         context.initialize()
         context.finalize()
-        context._write_file_footer = Mock(side_effect=Exception())
         context.finalize()
-        self.assertEquals(context._write_file_footer.call_count, 0)
+        self.assertEquals(context._write_root_tag_close.call_count, 1)
 
     def test_finalize_with_valid_checksum_type(self):
 
@@ -239,7 +266,7 @@ def test_finalize_with_valid_checksum_type(self):
         context = MetadataFileContext(path, checksum_type)
 
         context._open_metadata_file_handle()
-        context._write_file_header()
+        context._write_xml_header()
         context.finalize()
 
         expected_metadata_file_name = context.checksum + '-' + 'test.xml'
@@ -255,7 +282,7 @@ def test_finalize_error_on_footer(self, mock_logger):
         context._write_file_footer = Mock(side_effect=Exception('foo'))
 
         context._open_metadata_file_handle()
-        context._write_file_header()
+        context._write_xml_header()
 
         context.finalize()
 
@@ -269,7 +296,7 @@ def test_finalize_error_on_close(self, mock_logger):
         context._close_metadata_file_handle = Mock(side_effect=Exception('foo'))
 
         context._open_metadata_file_handle()
-        context._write_file_header()
+        context._write_xml_header()
 
         context.finalize()