SlurmGCP. Improve reservation_name parsing logic + tests

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/tests/common.py

@@ -29,11 +29,13 @@ class Placeholder:
 
 @dataclass
 class TstNodeset:
-    nodeset_name: str
+    nodeset_name: str = "cantor"
     node_count_static: int = 0
     node_count_dynamic_max: int = 0
     node_conf: dict[str, Any] = field(default_factory=dict)
     instance_template: Optional[str] = None
+    reservation_name: Optional[str] = ""
+    zone_policy_allow: Optional[list[str]] = field(default_factory=list)
 
 @dataclass
 class TstCfg:

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/tests/test_util.py

@@ -13,7 +13,8 @@
 # limitations under the License.
 
 import pytest
-import common # needed to import util
+from mock import Mock
+from common import TstNodeset, TstCfg # needed to import util
 import util
 from google.api_core.client_options import ClientOptions  # noqa: E402
 
@@ -130,3 +131,85 @@ def test_create_client_options(
     ud_mock.return_value = "googleapis.com"
     ep_mock.return_value = ep_ver
     assert util.create_client_options(api).__repr__() == expected.__repr__()
+
+
+
+@pytest.mark.parametrize(
+        "nodeset,err",
+        [
+            (TstNodeset(reservation_name="projects/x/reservations/y"), AssertionError), # no zones
+            (TstNodeset(
+                reservation_name="projects/x/reservations/y",
+                zone_policy_allow=["eine", "zwei"]), AssertionError), # multiples zones
+            (TstNodeset(
+                reservation_name="robin",
+                zone_policy_allow=["eine"]), ValueError), # invalid name
+            (TstNodeset(
+                reservation_name="projects/reservations/y",
+                zone_policy_allow=["eine"]), ValueError), # invalid name
+            (TstNodeset(
+                reservation_name="projects/x/zones/z/reservations/y",
+                zone_policy_allow=["eine"]), ValueError), # invalid name
+        ]
+)
+def test_nodeset_reservation_err(nodeset, err):
+    lkp = util.Lookup(TstCfg())
+    lkp._get_reservation = Mock()
+    with pytest.raises(err):
+        lkp.nodeset_reservation(nodeset)
+    lkp._get_reservation.assert_not_called()
+
+@pytest.mark.parametrize(
+        "nodeset,policies,expected",
+        [
+            (TstNodeset(), [], None), # no reservation
+            (TstNodeset(
+                reservation_name="projects/bobin/reservations/robin",
+                zone_policy_allow=["eine"]), 
+                [],
+                util.ReservationDetails(
+                    project="bobin",
+                    zone="eine",
+                    name="robin",
+                    policies=[],
+                    bulk_insert_name="projects/bobin/reservations/robin")),
+            (TstNodeset(
+                reservation_name="projects/bobin/reservations/robin",
+                zone_policy_allow=["eine"]), 
+                ["seven/wanders", "five/red/apples", "yum"],
+                util.ReservationDetails(
+                    project="bobin",
+                    zone="eine",
+                    name="robin",
+                    policies=["wanders", "apples", "yum"],
+                    bulk_insert_name="projects/bobin/reservations/robin")),
+            (TstNodeset(
+                reservation_name="projects/bobin/reservations/robin/reservationBlocks/cheese-brie-6",
+                zone_policy_allow=["eine"]), 
+                [],
+                util.ReservationDetails(
+                    project="bobin",
+                    zone="eine",
+                    name="robin",
+                    policies=[],
+                    reservation_block="cheese-brie-6",
+                    bulk_insert_name="projects/bobin/reservations/robin/reservationBlocks/cheese-brie-6")),
+
+        ])
+
+def test_nodeset_reservation_ok(nodeset, policies, expected):
+    lkp = util.Lookup(TstCfg())
+    lkp._get_reservation = Mock()
+
+    if not expected:
+        assert lkp.nodeset_reservation(nodeset) is None
+        lkp._get_reservation.assert_not_called()
+        return
+
+    lkp._get_reservation.return_value = {
+        "resourcePolicies": {i: p for i, p in enumerate(policies)},
+    }
+    assert lkp.nodeset_reservation(nodeset) == expected
+    lkp._get_reservation.assert_called_once_with(expected.project, expected.zone, expected.name)
+
+

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/scripts/util.py

@@ -1447,8 +1447,10 @@ def delete_node(self, nodename):
 class ReservationDetails:
     project: str
     zone: str
+    name: str
     policies: List[str] # names (not URLs) of resource policies
-    bulk_insert_name: str # name in format suitable for bulk insert (currently identical to user supplied name)
+    bulk_insert_name: str # name in format suitable for bulk insert (currently identical to user supplied name in long format)
+    reservation_block: Optional[str] = None
 
 class Lookup:
     """Wrapper class for cached data access"""
@@ -1754,13 +1756,13 @@ def nodeset_reservation(self, nodeset: object) -> Optional[ReservationDetails]:
         assert len(zones) == 1, "Only single zone is supported if using a reservation"
         zone = zones[0]
 
-        try:
-            _, project, _, name = nodeset.reservation_name.split("/")
-        except ValueError:
+        regex = re.compile(r'^projects/(?P<project>[^/]+)/reservations/(?P<reservation>[^/]+)(/reservationBlocks/(?P<block>[^/]+))?$')
+        if not (match := regex.match(nodeset.reservation_name)):
             raise ValueError(
-                f"Invalid reservation name: '{nodeset.reservation_name}', expected format is 'projects/PROJECT/reservations/NAME'"
+                f"Invalid reservation name: '{nodeset.reservation_name}', expected format is 'projects/PROJECT/reservations/NAME[/reservationBlocks/BLOCK]'"
             )
 
+        project, name, block = match.group("project", "reservation", "block")
         reservation = self._get_reservation(project, zone, name)
 
         # Converts policy URLs to names, e.g.:
@@ -1770,7 +1772,9 @@ def nodeset_reservation(self, nodeset: object) -> Optional[ReservationDetails]:
         return ReservationDetails(
             project=project,
             zone=zone,
+            name=name,
             policies=policies,
+            reservation_block=block,
             bulk_insert_name=nodeset.reservation_name)
 
     @lru_cache(maxsize=1)