diff --git a/src/pylorax/yumbase.py b/src/pylorax/yumbase.py
new file mode 100644
index 00000000..0646a2a9
--- /dev/null
+++ b/src/pylorax/yumbase.py
@@ -0,0 +1,137 @@
+# Copyright (C) 2009-2020 Red Hat, Inc.
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation; either version 2 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program. If not, see .
+#
+# Red Hat Author(s): Martin Gracik
+#
+# pylint: disable=bad-preconf-access
+
+
+import ConfigParser
+import logging
+import os
+import yum
+
+def get_yum_base_object(installroot, repositories, mirrorlists=None, repo_files=None,
+ tempdir="/var/tmp", proxy=None, excludepkgs=None,
+ sslverify=True, releasever=None):
+
+ def sanitize_repo(repo):
+ """Convert bare paths to file:/// URIs, and silently reject protocols unhandled by yum"""
+ if repo.startswith("/"):
+ return "file://{0}".format(repo)
+ elif any(repo.startswith(p) for p in ('http://', 'https://', 'ftp://', 'file://')):
+ return repo
+ else:
+ return None
+
+ if mirrorlists is None:
+ mirrorlists = []
+ if repo_files is None:
+ repo_files = []
+ if excludepkgs is None:
+ excludepkgs = []
+
+ # sanitize the repositories
+ repositories = map(sanitize_repo, repositories)
+ mirrorlists = map(sanitize_repo, mirrorlists)
+
+ # remove invalid repositories
+ repositories = filter(bool, repositories)
+ mirrorlists = filter(bool, mirrorlists)
+
+ cachedir = os.path.join(tempdir, "yum.cache")
+ if not os.path.isdir(cachedir):
+ os.mkdir(cachedir)
+
+ yumconf = os.path.join(tempdir, "yum.conf")
+ c = ConfigParser.ConfigParser()
+
+ # add the main section
+ section = "main"
+ data = {"cachedir": cachedir,
+ "keepcache": 0,
+ "gpgcheck": 0,
+ "plugins": 0,
+ "assumeyes": 1,
+ "reposdir": "",
+ "tsflags": "nodocs"}
+
+ if proxy:
+ data["proxy"] = proxy
+
+ if sslverify == False:
+ data["sslverify"] = "0"
+
+ if excludepkgs:
+ data["exclude"] = " ".join(excludepkgs)
+
+ c.add_section(section)
+ map(lambda (key, value): c.set(section, key, value), data.items())
+
+ # add the main repository - the first repository from list
+ # This list may be empty if using --repo to load .repo files
+ if repositories:
+ section = "lorax-repo"
+ data = {"name": "lorax repo",
+ "baseurl": repositories[0],
+ "enabled": 1}
+
+ c.add_section(section)
+ map(lambda (key, value): c.set(section, key, value), data.items())
+
+ # add the extra repositories
+ for n, extra in enumerate(repositories[1:], start=1):
+ section = "lorax-extra-repo-{0:d}".format(n)
+ data = {"name": "lorax extra repo {0:d}".format(n),
+ "baseurl": extra,
+ "enabled": 1}
+
+ c.add_section(section)
+ map(lambda (key, value): c.set(section, key, value), data.items())
+
+ # add the mirrorlists
+ for n, mirror in enumerate(mirrorlists, start=1):
+ section = "lorax-mirrorlist-{0:d}".format(n)
+ data = {"name": "lorax mirrorlist {0:d}".format(n),
+ "mirrorlist": mirror,
+ "enabled": 1 }
+
+ c.add_section(section)
+ map(lambda (key, value): c.set(section, key, value), data.items())
+
+ # write the yum configuration file
+ with open(yumconf, "w") as f:
+ c.write(f)
+
+ # create the yum base object
+ yb = yum.YumBase()
+
+ yb.preconf.fn = yumconf
+ yb.preconf.root = installroot
+ if releasever:
+ yb.preconf.releasever = releasever
+
+ # Turn on as much yum logging as we can
+ yb.preconf.debuglevel = 6
+ yb.preconf.errorlevel = 6
+ yb.logger.setLevel(logging.DEBUG)
+ yb.verbose_logger.setLevel(logging.DEBUG)
+
+ # Add .repo files from the cmdline
+ for fn in repo_files:
+ if os.path.exists(fn):
+ yb.getReposFromConfigFile(fn)
+
+ return yb
diff --git a/src/sbin/lorax b/src/sbin/lorax
index a83006d4..7ec62141 100755
--- a/src/sbin/lorax
+++ b/src/sbin/lorax
@@ -38,7 +38,6 @@ import os
import shutil
import tempfile
from optparse import OptionParser, OptionGroup
-import ConfigParser
import yum
# This is a bit of a hack to short circuit yum's internal logging
@@ -46,6 +45,7 @@ import yum
yum.logginglevels._added_handlers = True
import pylorax
from pylorax import DRACUT_DEFAULT, log_selinux_state
+from pylorax.yumbase import get_yum_base_object
VERSION = "{0}-{1}".format(os.path.basename(sys.argv[0]), pylorax.vernum)
@@ -307,119 +307,5 @@ def main(args):
os.close(dir_fd)
-def get_yum_base_object(installroot, repositories, mirrorlists=None, repo_files=None,
- tempdir="/var/tmp", proxy=None, excludepkgs=None,
- sslverify=True, releasever=None):
-
- def sanitize_repo(repo):
- """Convert bare paths to file:/// URIs, and silently reject protocols unhandled by yum"""
- if repo.startswith("/"):
- return "file://{0}".format(repo)
- elif any(repo.startswith(p) for p in ('http://', 'https://', 'ftp://', 'file://')):
- return repo
- else:
- return None
-
- if mirrorlists is None:
- mirrorlists = []
- if repo_files is None:
- repo_files = []
- if excludepkgs is None:
- excludepkgs = []
-
- # sanitize the repositories
- repositories = map(sanitize_repo, repositories)
- mirrorlists = map(sanitize_repo, mirrorlists)
-
- # remove invalid repositories
- repositories = filter(bool, repositories)
- mirrorlists = filter(bool, mirrorlists)
-
- cachedir = os.path.join(tempdir, "yum.cache")
- if not os.path.isdir(cachedir):
- os.mkdir(cachedir)
-
- yumconf = os.path.join(tempdir, "yum.conf")
- c = ConfigParser.ConfigParser()
-
- # add the main section
- section = "main"
- data = {"cachedir": cachedir,
- "keepcache": 0,
- "gpgcheck": 0,
- "plugins": 0,
- "assumeyes": 1,
- "reposdir": "",
- "tsflags": "nodocs"}
-
- if proxy:
- data["proxy"] = proxy
-
- if sslverify == False:
- data["sslverify"] = "0"
-
- if excludepkgs:
- data["exclude"] = " ".join(excludepkgs)
-
- c.add_section(section)
- map(lambda (key, value): c.set(section, key, value), data.items())
-
- # add the main repository - the first repository from list
- # This list may be empty if using --repo to load .repo files
- if repositories:
- section = "lorax-repo"
- data = {"name": "lorax repo",
- "baseurl": repositories[0],
- "enabled": 1}
-
- c.add_section(section)
- map(lambda (key, value): c.set(section, key, value), data.items())
-
- # add the extra repositories
- for n, extra in enumerate(repositories[1:], start=1):
- section = "lorax-extra-repo-{0:d}".format(n)
- data = {"name": "lorax extra repo {0:d}".format(n),
- "baseurl": extra,
- "enabled": 1}
-
- c.add_section(section)
- map(lambda (key, value): c.set(section, key, value), data.items())
-
- # add the mirrorlists
- for n, mirror in enumerate(mirrorlists, start=1):
- section = "lorax-mirrorlist-{0:d}".format(n)
- data = {"name": "lorax mirrorlist {0:d}".format(n),
- "mirrorlist": mirror,
- "enabled": 1 }
-
- c.add_section(section)
- map(lambda (key, value): c.set(section, key, value), data.items())
-
- # write the yum configuration file
- with open(yumconf, "w") as f:
- c.write(f)
-
- # create the yum base object
- yb = yum.YumBase()
-
- yb.preconf.fn = yumconf
- yb.preconf.root = installroot
- if releasever:
- yb.preconf.releasever = releasever
-
- # Turn on as much yum logging as we can
- yb.preconf.debuglevel = 6
- yb.preconf.errorlevel = 6
- yb.logger.setLevel(logging.DEBUG)
- yb.verbose_logger.setLevel(logging.DEBUG)
-
- # Add .repo files from the cmdline
- for fn in repo_files:
- if os.path.exists(fn):
- yb.getReposFromConfigFile(fn)
-
- return yb
-
-
if __name__ == "__main__":
main(sys.argv)
diff --git a/tests/pylorax/test_treebuilder.py b/tests/pylorax/test_treebuilder.py
new file mode 100644
index 00000000..e9550289
--- /dev/null
+++ b/tests/pylorax/test_treebuilder.py
@@ -0,0 +1,136 @@
+#
+# Copyright (C) 2020 Red Hat, Inc.
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation; either version 2 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program. If not, see .
+#
+from contextlib import contextmanager
+import os
+from rpmfluff import SimpleRpmBuild, SourceFile, expectedArch
+import shutil
+import tempfile
+import unittest
+
+from pylorax import ArchData, DataHolder
+from pylorax.dnfbase import get_dnf_base_object
+from pylorax.treebuilder import RuntimeBuilder
+
+# TODO Put these into a common test library location
+@contextmanager
+def in_tempdir(prefix='tmp'):
+ """Execute a block of code with chdir in a temporary location"""
+ oldcwd = os.getcwd()
+ tmpdir = tempfile.mkdtemp(prefix=prefix)
+ os.chdir(tmpdir)
+ try:
+ yield
+ finally:
+ os.chdir(oldcwd)
+ shutil.rmtree(tmpdir)
+
+def makeFakeRPM(repo_dir, name, epoch, version, release, files=None, provides=None):
+ """Make a fake rpm file in repo_dir"""
+ if provides is None:
+ provides = []
+ p = SimpleRpmBuild(name, version, release)
+ if epoch:
+ p.epoch = epoch
+ if not files:
+ p.add_simple_payload_file_random()
+ else:
+ # Make a number of fake file entries in the rpm
+ for f in files:
+ p.add_installed_file(
+ installPath = f,
+ sourceFile = SourceFile(os.path.basename(f), "THIS IS A FAKE FILE"))
+ for c in provides:
+ p.add_provides(c)
+ with in_tempdir("lorax-test-rpms."):
+ p.make()
+ rpmfile = p.get_built_rpm(expectedArch)
+ shutil.move(rpmfile, repo_dir)
+
+
+class InstallBrandingTestCase(unittest.TestCase):
+ def install_branding(self, repo_dir, variant=None):
+ """Run the _install_branding and return the names of the installed packages"""
+ with tempfile.TemporaryDirectory(prefix="lorax.test.") as root_dir:
+ dbo = get_dnf_base_object(root_dir, ["file://"+repo_dir], enablerepos=[], disablerepos=[])
+ self.assertTrue(dbo is not None)
+
+ product = DataHolder(name="Fedora", version="33", release="33",
+ variant=variant, bugurl="http://none", isfinal=True)
+ arch = ArchData(os.uname().machine)
+ rb = RuntimeBuilder(product, arch, dbo)
+ rb._install_branding()
+ dbo.resolve()
+ self.assertTrue(dbo.transaction is not None)
+
+ return sorted(p.name for p in dbo.transaction.install_set)
+
+ def test_no_pkgs(self):
+ """Test with a repo with no system-release packages"""
+ # No system-release packages
+ with tempfile.TemporaryDirectory(prefix="lorax.test.repo.") as repo_dir:
+ makeFakeRPM(repo_dir, "fake-milhouse", 0, "1.0.0", "1")
+ os.system("createrepo_c " + repo_dir)
+
+ pkgs = self.install_branding(repo_dir)
+ self.assertEqual(pkgs, [])
+
+ def test_generic_pkg(self):
+ """Test with a repo with only a generic-release package"""
+ # Only generic-release
+ with tempfile.TemporaryDirectory(prefix="lorax.test.repo.") as repo_dir:
+ makeFakeRPM(repo_dir, "generic-release", 0, "33", "1", ["/etc/system-release"], ["system-release"])
+ os.system("createrepo_c " + repo_dir)
+
+ pkgs = self.install_branding(repo_dir)
+ self.assertEqual(pkgs, [])
+
+ def test_two_pkgs(self):
+ """Test with a repo with generic-release, and a fedora-release package"""
+ # Two system-release packages
+ with tempfile.TemporaryDirectory(prefix="lorax.test.repo.") as repo_dir:
+ makeFakeRPM(repo_dir, "generic-release", 0, "33", "1", ["/etc/system-release"], ["system-release"])
+ makeFakeRPM(repo_dir, "fedora-release", 0, "33", "1", ["/etc/system-release"], ["system-release"])
+ makeFakeRPM(repo_dir, "fedora-logos", 0, "33", "1")
+ os.system("createrepo_c " + repo_dir)
+
+ pkgs = self.install_branding(repo_dir)
+ self.assertEqual(pkgs, ["fedora-logos", "fedora-release"])
+
+ # Test with a variant set, but not available
+ pkgs = self.install_branding(repo_dir, variant="workstation")
+ self.assertEqual(pkgs, ["fedora-logos", "fedora-release"])
+
+ def test_three_pkgs(self):
+ """Test with a repo with generic-release, fedora-release, fedora-release-workstation package"""
+ # Three system-release packages, one with a variant suffix
+ with tempfile.TemporaryDirectory(prefix="lorax.test.repo.") as repo_dir:
+ makeFakeRPM(repo_dir, "generic-release", 0, "33", "1", ["/etc/system-release"], ["system-release"])
+ makeFakeRPM(repo_dir, "fedora-release", 0, "33", "1", ["/etc/system-release"], ["system-release"])
+ makeFakeRPM(repo_dir, "fedora-logos", 0, "33", "1")
+ makeFakeRPM(repo_dir, "fedora-release-workstation", 0, "33", "1", ["/etc/system-release"], ["system-release"])
+ os.system("createrepo_c " + repo_dir)
+
+ pkgs = self.install_branding(repo_dir)
+ self.assertEqual(pkgs, ["fedora-logos", "fedora-release"])
+
+ # Test with a variant set
+ pkgs = self.install_branding(repo_dir, variant="workstation")
+ self.assertEqual(pkgs, ["fedora-logos", "fedora-release-workstation"])
+
+ # Test with a variant set, but not available
+ pkgs = self.install_branding(repo_dir, variant="server")
+ self.assertEqual(pkgs, ["fedora-logos", "fedora-release"])