diff --git a/pungi/phases/extra_isos.py b/pungi/phases/extra_isos.py
index bc547e61..6a688bab 100644
--- a/pungi/phases/extra_isos.py
+++ b/pungi/phases/extra_isos.py
@@ -22,7 +22,7 @@ import productmd.treeinfo
 from pungi import createiso
 from pungi.phases.base import ConfigGuardedPhase, PhaseBase, PhaseLoggerMixin
 from pungi.phases.createiso import (add_iso_to_metadata, copy_boot_images,
-                                    run_createiso_command)
+                                    run_createiso_command, load_and_tweak_treeinfo)
 from pungi.util import failable, get_format_substs, get_variant_data, get_volid
 from pungi.wrappers import iso
 from pungi.wrappers.scm import get_dir_from_scm, get_file_from_scm
@@ -196,8 +196,7 @@ def get_iso_contents(compose, variant, arch, include_variants, filename, bootabl
 
 
 def tweak_treeinfo(compose, include_variants, source_file, dest_file):
-    ti = productmd.treeinfo.TreeInfo()
-    ti.load(source_file)
+    ti = load_and_tweak_treeinfo(source_file)
     for variant_uid in include_variants:
         variant = compose.all_variants[variant_uid]
         var = productmd.treeinfo.Variant(ti)
diff --git a/tests/fixtures/extraiso-expected.treeinfo b/tests/fixtures/extraiso-expected.treeinfo
index 8922a54e..11538f17 100644
--- a/tests/fixtures/extraiso-expected.treeinfo
+++ b/tests/fixtures/extraiso-expected.treeinfo
@@ -1,6 +1,3 @@
-[checksums]
-images/boot.iso = sha256:fc8a4be604b6425746f12fa706116eb940f93358f036b8fbbe518b516cb6870c
-
 [general]
 ; WARNING.0 = This section provides compatibility with pre-productmd treeinfos.
 ; WARNING.1 = Read productmd documentation for details about new format.
@@ -20,12 +17,15 @@ type = productmd.treeinfo
 version = 1.2
 
 [images-x86_64]
-boot.iso = images/boot.iso
 
 [images-xen]
 initrd = images/pxeboot/initrd.img
 kernel = images/pxeboot/vmlinuz
 
+[media]
+discnum = 1
+totaldiscs = 1
+
 [release]
 name = Test
 short = T
diff --git a/tests/test_extra_isos_phase.py b/tests/test_extra_isos_phase.py
index 55a56ef7..e17202c1 100644
--- a/tests/test_extra_isos_phase.py
+++ b/tests/test_extra_isos_phase.py
@@ -630,13 +630,7 @@ class TweakTreeinfoTest(helpers.PungiTestCase):
         expected = os.path.join(helpers.FIXTURE_DIR, "extraiso-expected.treeinfo")
         extra_isos.tweak_treeinfo(compose, ["Client"], input, output)
 
-        with open(expected) as f:
-            expected = f.read()
-        with open(output) as f:
-            actual = f.read()
-
-        self.maxDiff = None
-        self.assertEqual(expected, actual)
+        self.assertFilesEqual(output, expected)
 
 
 if __name__ == '__main__':