From 7d265901bc9b8c54e3ca97c29eb586817a2b8c9c Mon Sep 17 00:00:00 2001
From: Andreas Eknes Lie <andrli@equinor.com>
Date: Thu, 16 Jan 2025 10:55:31 +0100
Subject: [PATCH] Update komodo matrix functions and tests

---
 komodo/matrix.py     | 46 ++++++++++++++++++++++++++++++--------------
 tests/test_matrix.py | 42 ++++++++++++++++++++++++++++++++++++++--
 2 files changed, 72 insertions(+), 16 deletions(-)

diff --git a/komodo/matrix.py b/komodo/matrix.py
index a1c3bee9f..1c2c2aa46 100644
--- a/komodo/matrix.py
+++ b/komodo/matrix.py
@@ -6,35 +6,53 @@
 
 import itertools
 import re
-from typing import Iterator, Sequence, Tuple
+from typing import Iterator, Sequence, Tuple, Optional, Dict
 
 
 def get_matrix(
     rhel_versions: Sequence[str],
     py_versions: Sequence[str],
-) -> Iterator[Tuple[str, str]]:
+    other_components: Optional[Dict[str, Sequence[str]]] = None,
+) -> Iterator[Tuple[str, str, str]]:
     """Return tuples of rhel version and Python version, representing the
     current release matrix.
     """
-    for product in itertools.product(rhel_versions, py_versions):
-        rh_ver, py_ver = product
-        yield (f"rhel{rh_ver}", f"py{str(py_ver).replace('.', '')}")
 
+    if other_components:
+        ((component_name, component_seq),) = (
+            other_components.items()
+        )  # only pick a single component
 
-def format_release(base: str, rhel_ver: str, py_ver: str) -> str:
+        for product in itertools.product(rhel_versions, py_versions, component_seq):
+            rh_ver, py_ver, other_ver = product
+            yield (
+                f"rhel{rh_ver}",
+                f"py{str(py_ver).replace('.', '')}",
+                f"{component_name}{other_ver}",
+            )
+    else:
+        for product in itertools.product(rhel_versions, py_versions):
+            rh_ver, py_ver = product
+            yield f"rhel{rh_ver}", f"py{str(py_ver).replace('.', '')}", None
+
+
+def format_release(
+    base: str, rhel_ver: str, py_ver: str, other_component: Optional[str] = None
+) -> str:
     """Format a base (e.g. a matrix file without the .yml suffix) such that it
     looks like a concrete release.
     """
-    return f"{base}-{py_ver}-{rhel_ver}"
+    return (
+        f"{base}-{py_ver}-{rhel_ver}-{other_component}"
+        if other_component
+        else f"{base}-{py_ver}-{rhel_ver}"
+    )
 
 
 def get_matrix_base(release_name: str) -> str:
-    """Return the base (e.g. matrix part of a concrete release). Should be the
-    inverse of format_release for actual, concrete matrix releases.
-    Hard-coded the suffix pattern '-py..-rhel.' or '-py...-rhel.'.
+    """Return the base (e.g. matrix part of a concrete release).
+    Match release name on -py[nno]-rhel[n] and delimit using that
     """
-    suffix = format_release("", "rhel[0-9]", r"py\d{2,3}")
-    if re.search(suffix, release_name):
-        return re.split(suffix, release_name)[0]
-    # no matrix suffix at all
+    if re.search(r"-py\d{2,3}-rhel[0-9]", release_name):
+        return release_name.split("-py")[0]
     return release_name
diff --git a/tests/test_matrix.py b/tests/test_matrix.py
index d8120fee8..1f6bb21e3 100644
--- a/tests/test_matrix.py
+++ b/tests/test_matrix.py
@@ -3,8 +3,15 @@
 from komodo import matrix
 
 
-def test_format_matrix():
-    assert matrix.format_release("base", "rhel6", "py27") == "base-py27-rhel6"
+@pytest.mark.parametrize(
+    ("base", "rhel", "python", "other", "expected"),
+    [
+        ("base", "rhel6", "py27", None, "base-py27-rhel6"),
+        ("base", "rhel6", "py27", "numpy1", "base-py27-rhel6-numpy1"),
+    ],
+)
+def test_format_matrix(base, rhel, python, other, expected):
+    assert matrix.format_release(base, rhel, python, other) == expected
 
 
 @pytest.mark.parametrize(
@@ -14,8 +21,39 @@ def test_format_matrix():
         ("1970.12.rc0-foo-py38-rhel7", "1970.12.rc0-foo"),
         ("1970.12.03", "1970.12.03"),
         (matrix.format_release("1970.12.04", "rhel7", "py38"), "1970.12.04"),
+        (matrix.format_release("1990.06.04", "rhel9", "py38", "numpy2"), "1990.06.04"),
         ("1970.12.05-rhel8-py27", "1970.12.05-rhel8-py27"),  # outside matrix
+        ("2025.02.00-py311-rhel9-numpy2", "2025.02.00"),
     ],
 )
 def test_get_matrix_base(test_input, expected):
     assert matrix.get_matrix_base(test_input) == expected
+
+
+@pytest.mark.parametrize(
+    ("rhel_ver", "py_ver", "other_ver", "expected_yield"),
+    [
+        (
+            ["8"],
+            ["38", "311"],
+            None,
+            [("rhel8", "py38", None), ("rhel8", "py311", None)],
+        ),
+        (
+            ["8", "9"],
+            ["311"],
+            {"numpy": ["1", "2"]},
+            [
+                ("rhel8", "py311", "numpy1"),
+                ("rhel8", "py311", "numpy2"),
+                ("rhel9", "py311", "numpy1"),
+                ("rhel9", "py311", "numpy2"),
+            ],
+        ),
+    ],
+)
+def test_get_matrix(rhel_ver, py_ver, other_ver, expected_yield):
+    yielded = []
+    for mat in matrix.get_matrix(rhel_ver, py_ver, other_ver):
+        yielded.append(mat)
+    assert yielded == expected_yield