Unify CapacityBlockReservationInfo with CapacityReservationInfo and a…

…dd unit tests

Signed-off-by: Enrico Usai <[email protected]>

src/aws/ec2.py

@@ -18,20 +18,36 @@ class CapacityReservationInfo:
     Data object wrapping the result of a describe-capacity-reservations call.
 
     {
-        "CapacityReservationId": "cr-abcdEXAMPLE9876ef ",
-        "EndDateType": "unlimited",
-        "AvailabilityZone": "eu-west-1a",
-        "InstanceMatchCriteria": "open",
-        "Tags": [],
-        "EphemeralStorage": false,
-        "CreateDate": "2019-08-07T11:34:19.000Z",
-        "AvailableInstanceCount": 3,
+        "CapacityReservationId": "cr-123456",
+        "OwnerId": "123",
+        "CapacityReservationArn": "arn:aws:ec2:us-east-2:123:capacity-reservation/cr-123456",
+        "AvailabilityZoneId": "use2-az1",
+        "InstanceType": "t3.large",
         "InstancePlatform": "Linux/UNIX",
-        "TotalInstanceCount": 3,
-        "State": "cancelled",
+        "AvailabilityZone": "eu-west-1a",
         "Tenancy": "default",
-        "EbsOptimized": true,
-        "InstanceType": "m5.large"
+        "TotalInstanceCount": 1,
+        "AvailableInstanceCount": 1,
+        "EbsOptimized": false,
+        "EphemeralStorage": false,
+        "State": "active",
+        "StartDate": "2023-11-15T11:30:00+00:00",
+        "EndDate": "2023-11-16T11:30:00+00:00",  # capacity-block only
+        "EndDateType": "limited",
+        "InstanceMatchCriteria": "targeted",
+        "CreateDate": "2023-10-25T20:40:13+00:00",
+        "Tags": [
+            {
+                "Key": "aws:ec2capacityreservation:incrementalRequestedQuantity",
+                "Value": "1"
+            },
+            {
+                "Key": "aws:ec2capacityreservation:capacityReservationType",
+                "Value": "capacity-block"
+            }
+        ],
+        "CapacityAllocations": [],
+        "ReservationType": "capacity-block"  # capacity-block only
     }
     """
 
@@ -46,41 +62,6 @@ def state(self):
         """Return the state of the Capacity Reservation."""
         return self.capacity_reservation_data.get("State")
 
-
-class CapacityBlockReservationInfo(CapacityReservationInfo):
-    """
-    Data object wrapping the result of a describe-capacity-reservations --capacity-type capacity-block call.
-
-    {   "CapacityReservationId": "cr-a1234567",
-        "EndDateType": "limited",
-        "ReservationType": "capacity-block",
-        "AvailabilityZone": "eu-east-2a",
-        "InstanceMatchCriteria":  "targeted",
-        "EphemeralStorage": false,
-        "CreateDate": "2023-07-29T14:22:45Z  ",
-        “StartDate": "2023-08-15T12:00:00Z",
-        “EndDate": "2023-08-19T12:00:00Z",
-        "AvailableInstanceCount": 0,
-        "InstancePlatform":  "Linux/UNIX",
-        "TotalInstanceCount": 16,
-        “State": "payment-pending",
-        "Tenancy":  "default",
-        "EbsOptimized": true,
-        "InstanceType": "p5.48xlarge“
-    }
-    """
-
-    def __init__(self, capacity_reservation_data):
-        super().__init__(capacity_reservation_data)
-
-    def start_date(self):
-        """Return the start date of the CB."""
-        return self.capacity_reservation_data.get("StartDate")
-
-    def end_date(self):
-        """Return the start date of the CB."""
-        return self.capacity_reservation_data.get("EndDate")
-
     def __eq__(self, other):
         return self.__dict__ == other.__dict__
 
@@ -92,19 +73,17 @@ def __init__(self, config=None):
         super().__init__("ec2", config=config)
 
     @AWSExceptionHandler.handle_client_exception
-    def describe_capacity_reservations(
-        self, capacity_reservation_ids: List[str], reservation_type=None
-    ) -> List[CapacityBlockReservationInfo]:
-        """Accept a space separated list of ids. Return a list of CapacityReservationInfo."""
+    def describe_capacity_reservations(self, capacity_reservation_ids: List[str]) -> List[CapacityReservationInfo]:
+        """Accept a space separated list of reservation ids. Return a list of CapacityReservationInfo."""
         result = []
         response = list(
             self._paginate_results(
                 self._client.describe_capacity_reservations,
                 CapacityReservationIds=capacity_reservation_ids,
-                ReservationType=reservation_type,
+                # ReservationType=reservation_type,  # not yet available
             )
         )
         for capacity_reservation in response:
-            result.append(CapacityBlockReservationInfo(capacity_reservation))
+            result.append(CapacityReservationInfo(capacity_reservation))
 
         return result

src/slurm_plugin/capacity_block_manager.py

@@ -25,7 +25,7 @@
 from slurm_plugin.slurm_resources import SlurmNode
 
 from aws.common import AWSClientError
-from aws.ec2 import CapacityBlockReservationInfo, Ec2Client
+from aws.ec2 import CapacityReservationInfo, Ec2Client
 
 logger = logging.getLogger(__name__)
 
@@ -66,8 +66,8 @@ def __init__(self, capacity_block_id, queue_name, compute_resource_name):
         self._capacity_block_reservation_info = None
         self._nodenames = []
 
-    def update_capacity_block_reservation_info(self, capacity_block_reservation_info: CapacityBlockReservationInfo):
-        """Update info from CapacityBlockReservationInfo."""
+    def update_capacity_block_reservation_info(self, capacity_block_reservation_info: CapacityReservationInfo):
+        """Update info from CapacityReservationInfo."""
         self._capacity_block_reservation_info = capacity_block_reservation_info
 
     def slurm_reservation_name(self):
@@ -314,7 +314,7 @@ def _update_capacity_blocks_info_from_ec2(self, capacity_blocks: Dict[str, Capac
         )
         try:
             capacity_block_reservations_info: List[
-                CapacityBlockReservationInfo
+                CapacityReservationInfo
             ] = self.ec2_client().describe_capacity_reservations(capacity_block_ids)
 
             for capacity_block_reservation_info in capacity_block_reservations_info:

tests/aws/test_ec2.py

@@ -0,0 +1,76 @@
+# Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+# with the License. A copy of the License is located at
+#
+# http://aws.amazon.com/apache2.0/
+#
+# or in the "LICENSE.txt" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+from collections import namedtuple
+
+import pytest
+from assertpy import assert_that
+
+from aws.common import AWSClientError
+from aws.ec2 import CapacityReservationInfo, Ec2Client
+
+MockedBoto3Request = namedtuple(
+    "MockedBoto3Request", ["method", "response", "expected_params", "generate_error", "error_code"]
+)
+# Set defaults for attributes of the namedtuple. Since fields with a default value must come after any fields without
+# a default, the defaults are applied to the rightmost parameters. In this case generate_error = False and
+# error_code = None
+MockedBoto3Request.__new__.__defaults__ = (False, None)
+
+
+@pytest.fixture()
+def boto3_stubber_path():
+    # we need to set the region in the environment because the Boto3ClientFactory requires it.
+    os.environ["AWS_DEFAULT_REGION"] = "us-east-1"
+    return "aws.common.boto3"
+
+
+FAKE_CAPACITY_BLOCK_ID = "cr-a1234567"
+FAKE_CAPACITY_BLOCK_INFO = {
+    "CapacityReservationId": FAKE_CAPACITY_BLOCK_ID,
+    "EndDateType": "limited",
+    # "ReservationType": "capacity-block",
+    "AvailabilityZone": "eu-east-2a",
+    "InstanceMatchCriteria": "targeted",
+    "EphemeralStorage": False,
+    "CreateDate": "2023-07-29T14:22:45Z  ",
+    "StartDate": "2023-08-15T12:00:00Z",
+    "EndDate": "2023-08-19T12:00:00Z",
+    "AvailableInstanceCount": 0,
+    "InstancePlatform": "Linux/UNIX",
+    "TotalInstanceCount": 16,
+    "State": "payment-pending",
+    "Tenancy": "default",
+    "EbsOptimized": True,
+    "InstanceType": "p5.48xlarge",
+}
+
+
+@pytest.mark.parametrize("generate_error", [True, False])
+def test_describe_capacity_reservations(boto3_stubber, generate_error):
+    """Verify that describe_capacity_reservations behaves as expected."""
+    dummy_message = "dummy error message"
+    mocked_requests = [
+        MockedBoto3Request(
+            method="describe_capacity_reservations",
+            expected_params={"CapacityReservationIds": [FAKE_CAPACITY_BLOCK_ID]},
+            response=dummy_message if generate_error else {"CapacityReservations": [FAKE_CAPACITY_BLOCK_INFO]},
+            generate_error=generate_error,
+            error_code=None,
+        )
+    ]
+    boto3_stubber("ec2", mocked_requests)
+    if generate_error:
+        with pytest.raises(AWSClientError, match=dummy_message):
+            Ec2Client().describe_capacity_reservations(capacity_reservation_ids=[FAKE_CAPACITY_BLOCK_ID])
+    else:
+        return_value = Ec2Client().describe_capacity_reservations(capacity_reservation_ids=[FAKE_CAPACITY_BLOCK_ID])
+        assert_that(return_value).is_equal_to([CapacityReservationInfo(FAKE_CAPACITY_BLOCK_INFO)])

tests/slurm_plugin/test_capacity_block_manager.py

@@ -24,7 +24,7 @@
 from slurm_plugin.slurm_resources import DynamicNode, SlurmReservation, StaticNode
 
 from aws.common import AWSClientError
-from aws.ec2 import CapacityBlockReservationInfo
+from aws.ec2 import CapacityReservationInfo
 
 FAKE_CAPACITY_BLOCK_ID = "cr-a1234567"
 FAKE_CAPACITY_BLOCK_INFO = {
@@ -58,7 +58,7 @@ class TestCapacityBlock:
         [("active", True), ("anything-else", False)],
     )
     def test_is_active(self, capacity_block, state, expected_output):
-        capacity_block_reservation_info = CapacityBlockReservationInfo({**FAKE_CAPACITY_BLOCK_INFO, "State": state})
+        capacity_block_reservation_info = CapacityReservationInfo({**FAKE_CAPACITY_BLOCK_INFO, "State": state})
         capacity_block.update_capacity_block_reservation_info(capacity_block_reservation_info)
         assert_that(capacity_block.is_active()).is_equal_to(expected_output)
 
@@ -147,11 +147,11 @@ def test_ec2_client(self, capacity_block_manager, mocker):
                     "cr-123456": CapacityBlock("cr-123456", "queue-cb", "compute-resource-cb"),
                 },
                 [
-                    CapacityBlockReservationInfo(
+                    CapacityReservationInfo(
                         {**FAKE_CAPACITY_BLOCK_INFO, "State": "pending", "CapacityReservationId": "cr-123456"}
                     ),
                     # add another id not in the capacity block to trigger an internal error, that does not stop the loop
-                    CapacityBlockReservationInfo(
+                    CapacityReservationInfo(
                         {**FAKE_CAPACITY_BLOCK_INFO, "State": "active", "CapacityReservationId": "cr-234567"}
                     ),
                 ],
@@ -189,13 +189,13 @@ def test_ec2_client(self, capacity_block_manager, mocker):
                     "cr-345678": CapacityBlock("cr-345678", "queue-cb3", "compute-resource-cb3"),
                 },
                 [
-                    CapacityBlockReservationInfo(
+                    CapacityReservationInfo(
                         {**FAKE_CAPACITY_BLOCK_INFO, "State": "pending", "CapacityReservationId": "cr-123456"}
                     ),
-                    CapacityBlockReservationInfo(
+                    CapacityReservationInfo(
                         {**FAKE_CAPACITY_BLOCK_INFO, "State": "active", "CapacityReservationId": "cr-234567"}
                     ),
-                    CapacityBlockReservationInfo(
+                    CapacityReservationInfo(
                         {**FAKE_CAPACITY_BLOCK_INFO, "State": "pending", "CapacityReservationId": "cr-345678"}
                     ),
                 ],
@@ -431,7 +431,7 @@ def test_update_slurm_reservation(
         caplog,
     ):
         caplog.set_level(logging.INFO)
-        capacity_block_reservation_info = CapacityBlockReservationInfo({**FAKE_CAPACITY_BLOCK_INFO, "State": state})
+        capacity_block_reservation_info = CapacityReservationInfo({**FAKE_CAPACITY_BLOCK_INFO, "State": state})
         capacity_block.update_capacity_block_reservation_info(capacity_block_reservation_info)
         capacity_block.add_nodename("node1")
         capacity_block.add_nodename("node2")
@@ -494,10 +494,10 @@ def test_update_slurm_reservation(
                     "cr-123456": CapacityBlock("id", "queue-cb", "compute-resource-cb"),
                 },
                 [
-                    CapacityBlockReservationInfo(
+                    CapacityReservationInfo(
                         {**FAKE_CAPACITY_BLOCK_INFO, "State": "active", "CapacityReservationId": "cr-123456"}
                     ),
-                    CapacityBlockReservationInfo(
+                    CapacityReservationInfo(
                         {**FAKE_CAPACITY_BLOCK_INFO, "State": "pending", "CapacityReservationId": "cr-234567"}
                     ),
                 ],
@@ -509,10 +509,10 @@ def test_update_slurm_reservation(
                     "cr-123456": CapacityBlock("id", "queue-cb", "compute-resource-cb"),
                 },
                 [
-                    CapacityBlockReservationInfo(
+                    CapacityReservationInfo(
                         {**FAKE_CAPACITY_BLOCK_INFO, "State": "active", "CapacityReservationId": "cr-123456"}
                     ),
-                    CapacityBlockReservationInfo(
+                    CapacityReservationInfo(
                         {**FAKE_CAPACITY_BLOCK_INFO, "State": "pending", "CapacityReservationId": "cr-234567"}
                     ),
                 ],
@@ -525,10 +525,10 @@ def test_update_slurm_reservation(
                     "cr-234567": CapacityBlock("cr-234567", "queue-cb", "compute-resource-cb"),
                 },
                 [
-                    CapacityBlockReservationInfo(
+                    CapacityReservationInfo(
                         {**FAKE_CAPACITY_BLOCK_INFO, "State": "active", "CapacityReservationId": "cr-123456"}
                     ),
-                    CapacityBlockReservationInfo(
+                    CapacityReservationInfo(
                         {**FAKE_CAPACITY_BLOCK_INFO, "State": "pending", "CapacityReservationId": "cr-234567"}
                     ),
                 ],
@@ -568,7 +568,7 @@ def test_update_capacity_blocks_info_from_ec2(
             assert_that(
                 capacity_block_manager._capacity_blocks.get("cr-123456")._capacity_block_reservation_info
             ).is_equal_to(
-                CapacityBlockReservationInfo(
+                CapacityReservationInfo(
                     {**FAKE_CAPACITY_BLOCK_INFO, "State": "active", "CapacityReservationId": "cr-123456"}
                 )
             )
@@ -580,14 +580,14 @@ def test_update_capacity_blocks_info_from_ec2(
                 assert_that(
                     capacity_block_manager._capacity_blocks.get("cr-123456")._capacity_block_reservation_info
                 ).is_equal_to(
-                    CapacityBlockReservationInfo(
+                    CapacityReservationInfo(
                         {**FAKE_CAPACITY_BLOCK_INFO, "State": "active", "CapacityReservationId": "cr-123456"}
                     )
                 )
                 assert_that(
                     capacity_block_manager._capacity_blocks.get("cr-234567")._capacity_block_reservation_info
                 ).is_equal_to(
-                    CapacityBlockReservationInfo(
+                    CapacityReservationInfo(
                         {**FAKE_CAPACITY_BLOCK_INFO, "State": "pending", "CapacityReservationId": "cr-234567"}
                     )
                 )