Merge pull request #261 from jimmysway/fix/256-pulling-outages

Integrate nerc-rates for dynamic outage information loading

Author: knikolla

requirements.txt

@@ -1,4 +1,4 @@
-git+https://github.com/CCI-MOC/nerc-rates@33701ed#egg=nerc_rates
+git+https://github.com/CCI-MOC/nerc-rates@5569bba#egg=nerc_rates
 boto3
 kubernetes
 openshift

src/coldfront_plugin_cloud/management/commands/calculate_storage_gb_hours.py

@@ -20,6 +20,12 @@
 
 _RATES = None
 
+RESOURCE_NAME_TO_NERC_SERVICE = {
+    "NERC": "stack",
+    "NERC-OCP": "ocp-prod",
+    "NERC-OCP-EDU": "academic",
+}
+
 
 def get_rates():
     # nerc-rates doesn't work with Python 3.9, which is what ColdFront is currently
@@ -142,13 +148,6 @@ def add_arguments(self, parser):
             action="store_true",
             help="Upload generated CSV invoice to S3 storage.",
         )
-        parser.add_argument(
-            "--excluded-time-ranges",
-            type=str,
-            default=None,
-            nargs="+",
-            help="List of time ranges excluded from billing, in ISO format.",
-        )
 
     @staticmethod
     def default_start_argument():
@@ -198,8 +197,24 @@ def upload_to_s3(s3_endpoint, s3_bucket, file_location, invoice_month):
         logger.info(f"Uploaded to {secondary_location}.")
 
     def handle(self, *args, **options):
+        def get_outages_for_service(resource_name: str):
+            """Get outages for a service from nerc-rates.
+
+            :param resource_name: Name of the resource to get outages for.
+            :return: List of excluded intervals or None.
+            """
+            service_name = RESOURCE_NAME_TO_NERC_SERVICE.get(resource_name)
+            if service_name:
+                return utils.load_outages_from_nerc_rates(
+                    options["start"], options["end"], service_name
+                )
+            return None
+
         def process_invoice_row(allocation, attrs, su_name, rate):
             """Calculate the value and write the bill using the writer."""
+            resource_name = allocation.resources.first().name
+            excluded_intervals_list = get_outages_for_service(resource_name)
+
             time = 0
             for attribute in attrs:
                 time += utils.calculate_quota_unit_hours(
@@ -234,13 +249,6 @@ def process_invoice_row(allocation, attrs, su_name, rate):
         logger.info(f"Processing invoices for {options['invoice_month']}.")
         logger.info(f"Interval {options['start'] - options['end']}.")
 
-        if options["excluded_time_ranges"]:
-            excluded_intervals_list = utils.load_excluded_intervals(
-                options["excluded_time_ranges"]
-            )
-        else:
-            excluded_intervals_list = None
-
         openstack_resources = Resource.objects.filter(
             resource_type=ResourceType.objects.get(name="OpenStack")
         )
@@ -285,15 +293,13 @@ def process_invoice_row(allocation, attrs, su_name, rate):
             csv_invoice_writer = csv.writer(
                 f, delimiter=",", quotechar="|", quoting=csv.QUOTE_MINIMAL
             )
-            # Write Headers
             csv_invoice_writer.writerow(InvoiceRow.get_headers())
 
             for allocation in openstack_allocations:
                 allocation_str = (
                     f'{allocation.pk} of project "{allocation.project.title}"'
                 )
-                msg = f"Starting billing for allocation {allocation_str}."
-                logger.debug(msg)
+                logger.debug(f"Starting billing for allocation {allocation_str}.")
 
                 process_invoice_row(
                     allocation,
@@ -306,8 +312,7 @@ def process_invoice_row(allocation, attrs, su_name, rate):
                 allocation_str = (
                     f'{allocation.pk} of project "{allocation.project.title}"'
                 )
-                msg = f"Starting billing for allocation {allocation_str}."
-                logger.debug(msg)
+                logger.debug(f"Starting billing for allocation {allocation_str}.")
 
                 process_invoice_row(
                     allocation,

src/coldfront_plugin_cloud/tests/unit/test_calculate_quota_unit_hours.py

@@ -1,6 +1,9 @@
+import csv
 import datetime
 import pytz
 import tempfile
+from decimal import Decimal
+from unittest.mock import Mock, patch
 
 import freezegun
 
@@ -16,7 +19,11 @@
 
 
 class TestCalculateAllocationQuotaHours(base.TestBase):
-    def test_new_allocation_quota(self):
+    @patch("coldfront_plugin_cloud.utils.load_outages_from_nerc_rates")
+    def test_new_allocation_quota(self, mock_load_outages):
+        """Test quota calculation with nerc-rates outages mocked."""
+        mock_load_outages.return_value = []
+
         self.resource = self.new_openshift_resource(
             name="",
         )
@@ -63,7 +70,7 @@ def test_new_allocation_quota(self):
                 "2020-03",
             )
 
-        # Let's test a complete CLI call including excluded time, while we're at it. This is not for testing
+        # Let's test a complete CLI call. This is not for testing
         # the validity but just the unerrored execution of the complete pipeline.
         # Tests that verify the correct output are further down in the test file.
         with tempfile.NamedTemporaryFile() as fp:
@@ -83,10 +90,12 @@ def test_new_allocation_quota(self):
                 "0.00001",
                 "--invoice-month",
                 "2020-03",
-                "--excluded-time-ranges",
-                "2020-03-02 00:00:00,2020-03-03 05:00:00",
             )
 
+        # Verify that load_outages_from_nerc_rates is not called when resource name
+        # doesn't match NERC service mapping
+        mock_load_outages.assert_not_called()
+
     def test_new_allocation_quota_expired(self):
         """Test that expiration doesn't affect invoicing."""
         self.resource = self.new_openshift_resource(
@@ -597,3 +606,65 @@ def test_load_excluded_intervals_invalid(self):
         ]
         with self.assertRaises(AssertionError):
             utils.load_excluded_intervals(invalid_interval)
+
+    @patch(
+        "coldfront_plugin_cloud.management.commands.calculate_storage_gb_hours.RESOURCE_NAME_TO_NERC_SERVICE",
+        {"TEST-RESOURCE": "test-service"},
+    )
+    @patch(
+        "coldfront_plugin_cloud.management.commands.calculate_storage_gb_hours.get_rates"
+    )
+    def test_nerc_outages_integration(self, mock_rates_loader):
+        """Test nerc-rates integration: get_outages_during called correctly and outages reduce billing."""
+        start = pytz.utc.localize(datetime.datetime(2020, 3, 1, 0, 0, 0))
+        end = pytz.utc.localize(datetime.datetime(2020, 3, 31, 0, 0, 0))
+        mock_outages = [
+            (
+                pytz.utc.localize(datetime.datetime(2020, 3, 10, 0, 0, 0)),
+                pytz.utc.localize(datetime.datetime(2020, 3, 12, 0, 0, 0)),
+            )
+        ]
+
+        mock_outages_data = Mock()
+        mock_outages_data.get_outages_during.return_value = mock_outages
+        mock_rates_loader.return_value.get_value_at.return_value = Decimal("0.001")
+
+        with patch.object(utils, "_OUTAGES_DATA", mock_outages_data):
+            with freezegun.freeze_time("2020-03-01"):
+                user = self.new_user()
+                project = self.new_project(pi=user)
+                resource = self.new_openstack_resource(name="TEST-RESOURCE")
+                allocation = self.new_allocation(project, resource, 100)
+                for attr, val in [
+                    (attributes.ALLOCATION_PROJECT_NAME, "test"),
+                    (attributes.ALLOCATION_PROJECT_ID, "123"),
+                    (attributes.QUOTA_VOLUMES_GB, 10),
+                ]:
+                    utils.set_attribute_on_allocation(allocation, attr, val)
+
+            with tempfile.NamedTemporaryFile(
+                mode="w", delete=False, suffix=".csv"
+            ) as fp:
+                output_file = fp.name
+
+            call_command(
+                "calculate_storage_gb_hours",
+                "--output",
+                output_file,
+                "--start",
+                "2020-03-01",
+                "--end",
+                "2020-03-31",
+                "--invoice-month",
+                "2020-03",
+            )
+
+            mock_outages_data.get_outages_during.assert_called_once_with(
+                start.isoformat(), end.isoformat(), "test-service"
+            )
+
+            with open(output_file, "r") as f:
+                rows = list(csv.DictReader(f))
+
+            self.assertEqual(len(rows), 1)
+            self.assertEqual(int(rows[0]["SU Hours (GBhr or SUhr)"]), 6720)

src/coldfront_plugin_cloud/utils.py

@@ -1,4 +1,5 @@
 import datetime
+import functools
 import math
 import pytz
 import re
@@ -14,6 +15,9 @@
 
 from coldfront_plugin_cloud import attributes
 
+# Load outages data once per program execution
+_OUTAGES_DATA = None
+
 
 def env_safe_name(name):
     return re.sub(r"[^A-Za-z0-9]", "_", str(name)).upper()
@@ -191,6 +195,12 @@ def calculate_quota_unit_hours(
 
 
 def load_excluded_intervals(excluded_interval_arglist):
+    """Parse excluded time ranges from command line arguments.
+
+    :param excluded_interval_arglist: List of time range strings in format "start,end".
+    :return: Sorted list of [start, end] datetime tuples.
+    """
+
     def interval_sort_key(e):
         return e[0]
 
@@ -223,6 +233,27 @@ def check_overlapping_intervals(excluded_intervals_list):
     return excluded_intervals_list
 
 
+@functools.cache
+def load_outages_from_nerc_rates(
+    start: datetime.datetime, end: datetime.datetime, affected_service: str
+) -> list[tuple[datetime.datetime, datetime.datetime]]:
+    """Load outage intervals from nerc-rates for a given time period and service.
+
+    :param start: Start time for outage search.
+    :param end: End time for outage search.
+    :param affected_service: Name of the affected service (e.g., "stack", "ocp-prod").
+    :return: List of [start, end] datetime tuples representing outages.
+    """
+    global _OUTAGES_DATA
+    if _OUTAGES_DATA is None:
+        from nerc_rates import outages
+
+        _OUTAGES_DATA = outages.load_from_url()
+    return _OUTAGES_DATA.get_outages_during(
+        start.isoformat(), end.isoformat(), affected_service
+    )
+
+
 def _clamp_time(time, min_time, max_time):
     if time < min_time:
         time = min_time