recursiveloader: Support specifying hashes in ctor

2 files changed, 109 insertions, 4 deletions

diff --git a/gemato/recursiveloader.py b/gemato/recursiveloader.py
index cb1b6d0..f403b7a 100644
--- a/gemato/recursiveloader.py
+++ b/gemato/recursiveloader.py
@@ -21,7 +21,8 @@ class ManifestRecursiveLoader(object):
 
     def __init__(self, top_manifest_path,
             verify_openpgp=True, openpgp_env=None,
-            sign_openpgp=None, openpgp_keyid=None):
+            sign_openpgp=None, openpgp_keyid=None,
+            hashes=None):
         """
         Instantiate the loader for a Manifest tree starting at top-level
         Manifest @top_manifest_path.
@@ -38,6 +39,11 @@ class ManifestRecursiveLoader(object):
         originally signed. @openpgp_keyid can be used to select the key.
 
         Sub-Manifests are never signed.
+
+        @hashes can be used to specify a default hash set
+        for the Manifest. If it is specified, they will be used for all
+        subsequent update*() calls that do not specify another set
+        of hashes explicitly.
         """
         self.root_directory = os.path.dirname(top_manifest_path)
         self.loaded_manifests = {}
@@ -45,6 +51,8 @@ class ManifestRecursiveLoader(object):
         self.openpgp_env = openpgp_env
         self.sign_openpgp = sign_openpgp
         self.openpgp_keyid = openpgp_keyid
+        self.hashes = hashes
+
         # TODO: allow catching OpenPGP exceptions somehow?
         m = self.load_manifest(os.path.basename(top_manifest_path))
         self.openpgp_signed = m.openpgp_signed
@@ -380,13 +388,20 @@ class ManifestRecursiveLoader(object):
         If the path does not exist, all Manifest entries for it will
         be removed except for OPTIONAL entries.
 
-        @hashes specifies the requested hash set. By default,
-        the existing hashes in the entry are updated. @hashes
-        must be specified when creating a new entry.
+        @hashes specifies the requested hash set. If specified,
+        it overrides the hash set used in Manifest. If None, the set
+        specified in ManifestLoader constructor is used. If that one
+        is None as well, the routine reuses the existing hash set
+        in the entry.
+
+        When creating a new entry, @hashes must be specified explicitly
+        either via the function or on construction.
         """
 
         had_entry = False
         manifests_to_update = set()
+        if hashes is None:
+            hashes = self.hashes
 
         self.load_manifests_for_path(path)
         for mpath, relpath, m in self._iter_manifests_for_path(path):
diff --git a/tests/test_recursiveloader.py b/tests/test_recursiveloader.py
index 53d1501..d017ef1 100644
--- a/tests/test_recursiveloader.py
+++ b/tests/test_recursiveloader.py
@@ -373,6 +373,40 @@ DATA test 0 MD5 d41d8cd98f00b204e9800998ecf8427e
         self.assertRaises(AssertionError,
                 m.update_entry_for_path, 'sub/stray')
 
+    def test_update_entry_for_path_hash_via_ctor(self):
+        m = gemato.recursiveloader.ManifestRecursiveLoader(
+            os.path.join(self.dir, 'Manifest'),
+            hashes=['SHA256', 'SHA512'])
+        m.update_entry_for_path('sub/stray')
+        self.assertListEqual(
+                sorted(m.find_path_entry('sub/stray').checksums),
+                ['SHA256', 'SHA512'])
+        # relevant Manifests should have been updated
+        with io.open(os.path.join(self.dir, 'sub/Manifest'),
+                     'r', encoding='utf8') as f:
+            self.assertNotEqual(f.read(), self.FILES['sub/Manifest'])
+        with io.open(os.path.join(self.dir, 'Manifest'),
+                     'r', encoding='utf8') as f:
+            self.assertNotEqual(f.read(), self.FILES['Manifest'])
+        m.assert_directory_verifies()
+
+    def test_update_entry_for_path_hash_via_ctor_and_override(self):
+        m = gemato.recursiveloader.ManifestRecursiveLoader(
+            os.path.join(self.dir, 'Manifest'),
+            hashes=['SHA256', 'SHA512'])
+        m.update_entry_for_path('sub/stray', hashes=['MD5'])
+        self.assertListEqual(
+                sorted(m.find_path_entry('sub/stray').checksums),
+                ['MD5'])
+        # relevant Manifests should have been updated
+        with io.open(os.path.join(self.dir, 'sub/Manifest'),
+                     'r', encoding='utf8') as f:
+            self.assertNotEqual(f.read(), self.FILES['sub/Manifest'])
+        with io.open(os.path.join(self.dir, 'Manifest'),
+                     'r', encoding='utf8') as f:
+            self.assertNotEqual(f.read(), self.FILES['Manifest'])
+        m.assert_directory_verifies()
+
 
 class MultipleManifestTest(TempDirTestCase):
     DIRS = ['sub']
@@ -465,6 +499,62 @@ TIMESTAMP 2017-01-01T01:01:01Z
             self.assertNotEqual(f.read(), self.FILES['Manifest'])
         m.assert_directory_verifies()
 
+    def test_update_entry_for_path_hashes_via_ctor(self):
+        m = gemato.recursiveloader.ManifestRecursiveLoader(
+            os.path.join(self.dir, 'Manifest'),
+            hashes=['SHA256', 'SHA512'])
+        m.update_entry_for_path('sub/foo')
+        # check for checksums
+        self.assertListEqual(
+                sorted(m.find_path_entry('sub/foo').checksums),
+                ['SHA256', 'SHA512'])
+        self.assertListEqual(
+                sorted(m.find_path_entry('sub/Manifest.a').checksums),
+                ['SHA256', 'SHA512'])
+        self.assertListEqual(
+                sorted(m.find_path_entry('sub/Manifest.b').checksums),
+                ['MD5'])
+        # relevant Manifests should have been updated
+        # but sub/Manifest.b should be left intact
+        with io.open(os.path.join(self.dir, 'sub/Manifest.a'),
+                     'r', encoding='utf8') as f:
+            self.assertNotEqual(f.read(), self.FILES['sub/Manifest.a'])
+        with io.open(os.path.join(self.dir, 'sub/Manifest.b'),
+                     'r', encoding='utf8') as f:
+            self.assertEqual(f.read(), self.FILES['sub/Manifest.b'])
+        with io.open(os.path.join(self.dir, 'Manifest'),
+                     'r', encoding='utf8') as f:
+            self.assertNotEqual(f.read(), self.FILES['Manifest'])
+        m.assert_directory_verifies()
+
+    def test_update_entry_for_path_hashes_via_ctor_and_override(self):
+        m = gemato.recursiveloader.ManifestRecursiveLoader(
+            os.path.join(self.dir, 'Manifest'),
+            hashes=['SHA256', 'SHA512'])
+        m.update_entry_for_path('sub/foo', hashes=['SHA1'])
+        # check for checksums
+        self.assertListEqual(
+                sorted(m.find_path_entry('sub/foo').checksums),
+                ['SHA1'])
+        self.assertListEqual(
+                sorted(m.find_path_entry('sub/Manifest.a').checksums),
+                ['SHA1'])
+        self.assertListEqual(
+                sorted(m.find_path_entry('sub/Manifest.b').checksums),
+                ['MD5'])
+        # relevant Manifests should have been updated
+        # but sub/Manifest.b should be left intact
+        with io.open(os.path.join(self.dir, 'sub/Manifest.a'),
+                     'r', encoding='utf8') as f:
+            self.assertNotEqual(f.read(), self.FILES['sub/Manifest.a'])
+        with io.open(os.path.join(self.dir, 'sub/Manifest.b'),
+                     'r', encoding='utf8') as f:
+            self.assertEqual(f.read(), self.FILES['sub/Manifest.b'])
+        with io.open(os.path.join(self.dir, 'Manifest'),
+                     'r', encoding='utf8') as f:
+            self.assertNotEqual(f.read(), self.FILES['Manifest'])
+        m.assert_directory_verifies()
+
 
 class MultipleTopLevelManifestTest(TempDirTestCase):
     FILES = {