lorax: Move get_yum_base_object into pylorax.yumbase

This will make it easier to run tests.

Related: rhbz#1472622

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 <http://www.gnu.org/licenses/>.
+#
+# Red Hat Author(s):  Martin Gracik <mgracik@redhat.com>
+#
+# 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

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)

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 <http://www.gnu.org/licenses/>.
+#
+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"])