createiso: Move code for tweaking treeinfo into a function

The function loads existing treeinfo, removes reference to boot.iso and
adds [media] section. This is the basic tweak that should happen for all
ISOs. Additional changes depend on the actual content.

Signed-off-by: Lubomír Sedlář <lsedlar@redhat.com>

pungi/phases/createiso.py

@@ -403,21 +403,7 @@ def prepare_iso(compose, arch, variant, disc_num=1, disc_count=None, split_iso_d
 
     # modify treeinfo
     ti_path = os.path.join(tree_dir, ".treeinfo")
-    ti = productmd.treeinfo.TreeInfo()
-    ti.load(ti_path)
-    ti.media.totaldiscs = disc_count or 1
-    ti.media.discnum = disc_num
-
-    # remove boot.iso from all sections
-    paths = set()
-    for platform in ti.images.images:
-        if "boot.iso" in ti.images.images[platform]:
-            paths.add(ti.images.images[platform].pop("boot.iso"))
-
-    # remove boot.iso from checksums
-    for i in paths:
-        if i in ti.checksums.checksums.keys():
-            del ti.checksums.checksums[i]
+    ti = load_and_tweak_treeinfo(ti_path, disc_num, disc_count)
 
     copy_boot_images(tree_dir, iso_dir)
 
@@ -447,8 +433,16 @@ def prepare_iso(compose, arch, variant, disc_num=1, disc_count=None, split_iso_d
             run("cp -a %s/repodata %s/" % (shlex_quote(tree_dir), shlex_quote(iso_dir)))
             with open(file_list, "w") as f:
                 f.write("\n".join(file_list_content))
-            cmd = repo.get_createrepo_cmd(tree_dir, update=True, database=True, skip_stat=True, pkglist=file_list, outputdir=iso_dir,
-                                          workers=compose.conf["createrepo_num_workers"], checksum=createrepo_checksum)
+            cmd = repo.get_createrepo_cmd(
+                tree_dir,
+                update=True,
+                database=True,
+                skip_stat=True,
+                pkglist=file_list,
+                outputdir=iso_dir,
+                workers=compose.conf["createrepo_num_workers"],
+                checksum=createrepo_checksum,
+            )
             run(cmd)
             # add repodata/repomd.xml back to checksums
             ti.checksums.add("repodata/repomd.xml", "sha256", root_dir=iso_dir)
@@ -482,6 +476,29 @@ def prepare_iso(compose, arch, variant, disc_num=1, disc_count=None, split_iso_d
     return gp
 
 
+def load_and_tweak_treeinfo(ti_path, disc_num=1, disc_count=1):
+    """Treeinfo on the media should not contain any reference to boot.iso and
+    it should also have a valid [media] section.
+    """
+    ti = productmd.treeinfo.TreeInfo()
+    ti.load(ti_path)
+    ti.media.totaldiscs = disc_count or 1
+    ti.media.discnum = disc_num
+
+    # remove boot.iso from all sections
+    paths = set()
+    for platform in ti.images.images:
+        if "boot.iso" in ti.images.images[platform]:
+            paths.add(ti.images.images[platform].pop("boot.iso"))
+
+    # remove boot.iso from checksums
+    for i in paths:
+        if i in ti.checksums.checksums.keys():
+            del ti.checksums.checksums[i]
+
+    return ti
+
+
 def copy_boot_images(src, dest):
     """When mkisofs is called it tries to modify isolinux/isolinux.bin and
     images/boot.img. Therefore we need to make copies of them.

tests/fixtures/expected-treeinfo

@@ -0,0 +1,59 @@
+[checksums]
+images/efiboot.img = sha256:e575ffee875b0135d378bdbd3277dea06377747b5daa689b14d684654d883817
+images/install.img = sha256:8106f4ce2ac973dcab7d103dcd02e7549510e7940342866931fe034661d67915
+images/pxeboot/initrd.img = sha256:cd6150fa7d9646febbc9535bdca14cddcfffc454b43c5d352f6d459d81792138
+images/pxeboot/vmlinuz = sha256:feb72d4112b4329db3fc12b6afcb27a6f96da00bb114b3b85527dfd11bdfc84f
+
+[general]
+; WARNING.0 = This section provides compatibility with pre-productmd treeinfos.
+; WARNING.1 = Read productmd documentation for details about new format.
+arch = x86_64
+family = Test Product
+name = Test Product 1.0
+packagedir = Packages
+platforms = x86_64,xen
+repository = .
+timestamp = 1539226055
+variant = Server
+variants = Server
+version = 1.0
+
+[header]
+type = productmd.treeinfo
+version = 1.2
+
+[images-x86_64]
+efiboot.img = images/efiboot.img
+initrd = images/pxeboot/initrd.img
+kernel = images/pxeboot/vmlinuz
+
+[images-xen]
+initrd = images/pxeboot/initrd.img
+kernel = images/pxeboot/vmlinuz
+
+[media]
+discnum = 1
+totaldiscs = 1
+
+[release]
+name = Test Product
+short = TP
+version = 1.0
+
+[stage2]
+mainimage = images/install.img
+
+[tree]
+arch = x86_64
+build_timestamp = 1539226055
+platforms = x86_64,xen
+variants = Server
+
+[variant-Server]
+id = Server
+name = Server
+packages = Packages
+repository = .
+type = variant
+uid = Server
+

tests/fixtures/original-treeinfo

@@ -0,0 +1,56 @@
+[checksums]
+images/boot.iso = sha256:5eaa52bee64c7823fe0a282247e41aeac51865ed0a2b0c7a1b85c1b8d1b835de
+images/efiboot.img = sha256:e575ffee875b0135d378bdbd3277dea06377747b5daa689b14d684654d883817
+images/install.img = sha256:8106f4ce2ac973dcab7d103dcd02e7549510e7940342866931fe034661d67915
+images/pxeboot/initrd.img = sha256:cd6150fa7d9646febbc9535bdca14cddcfffc454b43c5d352f6d459d81792138
+images/pxeboot/vmlinuz = sha256:feb72d4112b4329db3fc12b6afcb27a6f96da00bb114b3b85527dfd11bdfc84f
+
+[general]
+; WARNING.0 = This section provides compatibility with pre-productmd treeinfos.
+; WARNING.1 = Read productmd documentation for details about new format.
+arch = x86_64
+family = Test Product
+name = Test Product 1.0
+packagedir = Packages
+platforms = x86_64,xen
+repository = .
+timestamp = 1539226055
+variant = Server
+variants = Server
+version = 1.0
+
+[header]
+type = productmd.treeinfo
+version = 1.2
+
+[images-x86_64]
+boot.iso = images/boot.iso
+efiboot.img = images/efiboot.img
+initrd = images/pxeboot/initrd.img
+kernel = images/pxeboot/vmlinuz
+
+[images-xen]
+initrd = images/pxeboot/initrd.img
+kernel = images/pxeboot/vmlinuz
+
+[release]
+name = Test Product
+short = TP
+version = 1.0
+
+[stage2]
+mainimage = images/install.img
+
+[tree]
+arch = x86_64
+build_timestamp = 1539226055
+platforms = x86_64,xen
+variants = Server
+
+[variant-Server]
+id = Server
+name = Server
+packages = Packages
+repository = .
+type = variant
+uid = Server

tests/helpers.py

@@ -1,24 +1,39 @@
 # -*- coding: utf-8 -*-
 
-import mock
+import difflib
+import errno
+import imp
 import os
+import shutil
+import tempfile
+from collections import defaultdict
+
+import mock
+import six
+from kobo.rpmlib import parse_nvr
+
 try:
     import unittest2 as unittest
 except ImportError:
     import unittest
-import tempfile
-import shutil
-import errno
-import imp
-import six
-from kobo.rpmlib import parse_nvr
-from collections import defaultdict
 
 from pungi.util import get_arch_variant_data
 from pungi import paths, checks, Modulemd
 
 
-class PungiTestCase(unittest.TestCase):
+class BaseTestCase(unittest.TestCase):
+
+    def assertFilesEqual(self, fn1, fn2):
+        with open(fn1, 'rb') as f1:
+            lines1 = f1.read().decode('utf-8').splitlines()
+        with open(fn2, 'rb') as f2:
+            lines2 = f2.read().decode('utf-8').splitlines()
+        diff = '\n'.join(difflib.unified_diff(lines1, lines2,
+                                              fromfile='EXPECTED', tofile='ACTUAL'))
+        self.assertEqual(diff, '', 'Files differ:\n' + diff)
+
+
+class PungiTestCase(BaseTestCase):
     def setUp(self):
         self.topdir = tempfile.mkdtemp()
 

tests/test_comps_wrapper.py

@@ -5,7 +5,6 @@ try:
 except ImportError:
     import unittest
 import tempfile
-import difflib
 
 import os
 import sys
@@ -13,7 +12,7 @@ import sys
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
 
 from pungi.wrappers.comps import CompsWrapper, CompsFilter, CompsValidationError
-from tests.helpers import FIXTURE_DIR
+from tests.helpers import BaseTestCase, FIXTURE_DIR
 
 COMPS_FILE = os.path.join(FIXTURE_DIR, 'comps.xml')
 COMPS_FORMATTED_FILE = os.path.join(FIXTURE_DIR, 'comps-formatted.xml')
@@ -23,19 +22,10 @@ COMPS_FILE_WITH_TYPO = os.path.join(FIXTURE_DIR, 'comps-typo.xml')
 COMPS_FILE_WITH_WHITESPACE = os.path.join(FIXTURE_DIR, 'comps-ws.xml')
 
 
-class CompsWrapperTest(unittest.TestCase):
+class CompsWrapperTest(BaseTestCase):
     def setUp(self):
         self.file = tempfile.NamedTemporaryFile(prefix='comps-wrapper-test-')
 
-    def assertFilesEqual(self, fn1, fn2):
-        with open(fn1, 'rb') as f1:
-            lines1 = f1.read().decode('utf-8').splitlines()
-        with open(fn2, 'rb') as f2:
-            lines2 = f2.read().decode('utf-8').splitlines()
-        diff = '\n'.join(difflib.unified_diff(lines1, lines2,
-                                              fromfile='EXPECTED', tofile='ACTUAL'))
-        self.assertEqual(diff, '', 'Files differ:\n' + diff)
-
     def test_get_groups(self):
         comps = CompsWrapper(COMPS_FILE)
         self.assertEqual(

tests/test_createiso_phase.py

@@ -945,5 +945,17 @@ class BreakHardlinksTest(helpers.PungiTestCase):
         self.assertTrue(os.path.exists(expected))
 
 
+class TweakTreeinfo(helpers.PungiTestCase):
+    def test_tweaking(self):
+        input = os.path.join(helpers.FIXTURE_DIR, "original-treeinfo")
+        expected = os.path.join(helpers.FIXTURE_DIR, "expected-treeinfo")
+        output = os.path.join(self.topdir, "output")
+
+        ti = createiso.load_and_tweak_treeinfo(input)
+        ti.dump(output)
+
+        self.assertFilesEqual(output, expected)
+
+
 if __name__ == '__main__':
     unittest.main()