diff --git a/requirements.txt b/requirements.txt index 181a262..0a69b9b 100644 --- a/requirements.txt +++ b/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 diff --git a/src/coldfront_plugin_cloud/management/commands/calculate_storage_gb_hours.py b/src/coldfront_plugin_cloud/management/commands/calculate_storage_gb_hours.py index 170bcd0..f700738 100644 --- a/src/coldfront_plugin_cloud/management/commands/calculate_storage_gb_hours.py +++ b/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, diff --git a/src/coldfront_plugin_cloud/tests/unit/test_calculate_quota_unit_hours.py b/src/coldfront_plugin_cloud/tests/unit/test_calculate_quota_unit_hours.py index f704cde..c232881 100644 --- a/src/coldfront_plugin_cloud/tests/unit/test_calculate_quota_unit_hours.py +++ b/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) diff --git a/src/coldfront_plugin_cloud/utils.py b/src/coldfront_plugin_cloud/utils.py index d311481..6733fc4 100644 --- a/src/coldfront_plugin_cloud/utils.py +++ b/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