Add branch prefix [RHELDST-36928]

tests/ubiconfig/test_ubi.py

@@ -465,6 +465,44 @@ def test_pre_load(mocked_get_branches, mocked_session, files_branch_map, caplog)
     )
 
 
+@patch("requests.Session")
+@patch("ubiconfig._impl.loaders._GitlabLoader._get_branches")
+def test_pre_load_with_branch_prefix(mocked_get_branches, mocked_session, caplog):
+    branch_sha1 = OrderedDict(
+        [
+            ("ubi7.1", "2189cbc2e447f796fe354f8d784d76b0a2620248"),
+            ("ubi7", "c99cb8d7dae2e78e8cc7e720d3f950d1c5a0b51f"),
+            ("test-prefix8.1", "26d24af7859df3c4d361bd33cd57984d03abe206"),
+            ("test-prefix8", "117d4af7859df3c4d361bd33cdab6784d03abe104"),
+        ]
+    )
+    mocked_get_branches.return_value = branch_sha1
+    headers = {"Content-Length": "629", "X-Total-Pages": "1", "X-Per-Page": "20"}
+    mocked_session.return_value.request.return_value.headers = headers
+
+    file_list_ubi7_1 = [
+        {"name": "rhel-7-server.yaml", "path": "rhel-7-server.yaml"},
+    ]
+    file_list_ubi7 = [
+        {"name": "rhel-7-server.yaml", "path": "rhel-7-server.yaml"},
+    ]
+
+    mocked_session.return_value.request.return_value.json.side_effect = [
+        file_list_ubi7_1,
+        file_list_ubi7,
+    ]
+
+    loader = ubi.get_loader(branch_prefix="ubi")
+
+    actual_files_branch_map = loader._files_branch_map
+
+    assert "rhel-7-server.yaml" in actual_files_branch_map
+    assert len(actual_files_branch_map["rhel-7-server.yaml"]) == 2
+    assert actual_files_branch_map["rhel-7-server.yaml"][0][0] == "ubi7.1"
+    assert actual_files_branch_map["rhel-7-server.yaml"][1][0] == "ubi7"
+    assert "Skipping branch test-prefix" in caplog.text
+
+
 def test_ubi_config(ubi7_1_config_file1):
     config_dict = yaml.safe_load(ubi7_1_config_file1)
     config = UbiConfig.load_from_dict(config_dict, "rhel-atomic-host.yaml", "7.1")

ubiconfig/_impl/loaders/gitlab.py

@@ -23,12 +23,13 @@
 class GitlabLoader(Loader):
     """Load configuration from a remote repo on gitlab."""
 
-    def __init__(self, url):
+    def __init__(self, url, branch_prefix=None):
         """
         :param url: gitlab repo url in form of `https://<host>/<repo>`
         """
         self._url = url
         self._session = None
+        self._branch_prefix = branch_prefix
         self._repo_api = RepoApi(self._url.rstrip("/"))
         self._branches = self._get_branches()
         self._files_branch_map = self._pre_load()
@@ -85,7 +86,6 @@ def load(self, file_name, version=None):
         match = re.match(PREFIX_VERSION_RE, version)  # type: ignore
         prefix = match.group("prefix")  # type: ignore
         default_version = match.group("default_version")  # type: ignore
-
         default_branch = f"{prefix}{default_version}"
 
         sha1 = None
@@ -141,12 +141,22 @@ def _pre_load(self):
         LOG.debug("Loading config files from all branches")
 
         for branch, sha1 in self._branches.items():
-            if not re.match(PREFIX_VERSION_RE, branch):
+            branch_regex = re.match(PREFIX_VERSION_RE, branch)
+            if not branch_regex:
                 LOG.warning(
                     "Skipping branch %s (name does not match with required format)",
                     branch,
                 )
                 continue
+            prefix = branch_regex.group("prefix")
+            if self._branch_prefix and self._branch_prefix != prefix:
+                LOG.warning(
+                    "Skipping branch %s \
+                        (branch does not match with cdn definition branch %s)",
+                    prefix,
+                    self._branch_prefix,
+                )
+                continue
             page = 1
             while True:
                 file_list_api = self._repo_api.get_file_list_api(branch=sha1, page=page)

ubiconfig/ubi.py

@@ -13,7 +13,7 @@ class LoaderError(RuntimeError):
     pass
 
 
-def get_loader(source=None):
+def get_loader(source=None, branch_prefix=None):
     """Get a Loader instance which is used to load configurations.
 
     ``source`` should be provided as one of the following:
@@ -60,7 +60,7 @@ def get_loader(source=None):
     parsed = urlparse(source)
     if parsed.netloc:
         # It's a URL, use the gitlab loader
-        return _GitlabLoader(source)
+        return _GitlabLoader(source, branch_prefix)
 
     # It should be a local path
     if not os.path.isdir(source):