Split up test_discount_processor.py into tests for the discount processor and pi-su credit processor

process_report/tests/unit/processors/test_discount_processor.py

@@ -3,7 +3,7 @@
 import pandas
 
 from process_report.invoices import invoice
-from process_report.processors.pi_su_credit_processor import PISUCreditProcessor
+from process_report.processors.discount_processor import DiscountProcessor
 
 
 class TestDiscountProcessor(TestCase):
@@ -52,13 +52,21 @@ def test_preexisting_credit(
             balance=[90],
         )
 
-        processor = PISUCreditProcessor(
+        processor = DiscountProcessor(
             invoice_month="2024-06",
             data=invoice_data,
             name="test",
-            pi_su_mapping={"PI": ["Openstack Storage"]},
         )
-        processor.process()
+        processor.apply_flat_discount(
+            invoice=invoice_data,
+            pi_projects=invoice_data,
+            pi_balance_field=invoice.PI_BALANCE_FIELD,
+            discount_amount=invoice_data[invoice.COST_FIELD].sum(),
+            discount_field=invoice.CREDIT_FIELD,
+            balance_field=invoice.BALANCE_FIELD,
+            code_field=invoice.CREDIT_CODE_FIELD,
+            discount_code="0005",
+        )
         output_invoice = processor.data
 
         expected_invoice = self._get_test_invoice(
@@ -72,67 +80,3 @@ def test_preexisting_credit(
         )
 
         assert expected_invoice.equals(output_invoice)
-
-    def test_one_eligible_project_only(self):
-        """If PI has multiple projects but only one is eligible, only that project should get the credit"""
-        invoice_data = self._get_test_invoice(
-            pi=["PI", "PI"],
-            costs=[50, 75],
-            su_type=["Openstack Storage", "HPC"],
-        )
-
-        processor = PISUCreditProcessor(
-            invoice_month="2024-06",
-            data=invoice_data,
-            name="test",
-            pi_su_mapping={
-                "PI": ["Openstack Storage"]
-            },  # Only Openstack Storage SU type is eligible for credit, not HPC
-        )
-        processor.process()
-        output_invoice = processor.data
-
-        expected_invoice = self._get_test_invoice(
-            pi=["PI", "PI"],
-            costs=[50, 75],
-            su_type=["Openstack Storage", "HPC"],
-            credit=[50, None],
-            credit_code=["0005", None],
-            pi_balance=[0, 75],
-            balance=[0, 75],
-        )
-
-        expected_invoice = expected_invoice.astype(output_invoice.dtypes)
-        assert expected_invoice.equals(output_invoice)
-
-    def test_all_eligible_projects(self):
-        """PI has multiple projects and all are eligible, so all should get the credit"""
-        invoice_data = self._get_test_invoice(
-            pi=["PI", "PI"],
-            costs=[40, 60],
-            su_type=["Openstack Storage", "Openstack Compute"],
-        )
-
-        processor = PISUCreditProcessor(
-            invoice_month="2024-06",
-            data=invoice_data,
-            name="test",
-            pi_su_mapping={
-                "PI": ["Openstack Storage", "Openstack Compute"]
-            },  # Both SU types are eligible for credit
-        )
-        processor.process()
-        output_invoice = processor.data
-
-        expected_invoice = self._get_test_invoice(
-            pi=["PI", "PI"],
-            costs=[40, 60],
-            su_type=["Openstack Storage", "Openstack Compute"],
-            credit=[40, 60],
-            credit_code=["0005", "0005"],
-            pi_balance=[0, 0],
-            balance=[0, 0],
-        )
-
-        expected_invoice = expected_invoice.astype(output_invoice.dtypes)
-        assert expected_invoice.equals(output_invoice)

process_report/tests/unit/processors/test_pi_su_credit_processor.py

@@ -0,0 +1,103 @@
+from unittest import TestCase
+
+import pandas
+
+from process_report.invoices import invoice
+from process_report.tests import util as test_utils
+
+
+class TestPISUCreditProcessor(TestCase):
+    def _get_test_invoice(
+        self,
+        pi,
+        costs,
+        su_type,
+        credit=None,
+        credit_code=None,
+        pi_balance=None,
+        balance=None,
+    ):
+        if credit is None:
+            credit = [None for _ in range(len(pi))]
+        if credit_code is None:
+            credit_code = [None for _ in range(len(pi))]
+        if pi_balance is None:
+            pi_balance = costs
+        if balance is None:
+            balance = costs
+
+        return pandas.DataFrame(
+            {
+                invoice.PI_FIELD: pi,
+                invoice.SU_TYPE_FIELD: su_type,
+                invoice.COST_FIELD: costs,
+                invoice.CREDIT_FIELD: credit,
+                invoice.CREDIT_CODE_FIELD: credit_code,
+                invoice.PI_BALANCE_FIELD: pi_balance,
+                invoice.BALANCE_FIELD: balance,
+            }
+        )
+
+    def test_one_eligible_project_only(self):
+        """If PI has multiple projects but only one is eligible, only that project should get the credit"""
+
+        invoice_data = self._get_test_invoice(
+            pi=["PI", "PI"],
+            costs=[50, 75],
+            su_type=["Openstack Storage", "HPC"],
+        )
+
+        processor = test_utils.new_pi_su_credit_processor(
+            invoice_month="2024-06",
+            data=invoice_data,
+            pi_su_mapping={
+                "PI": ["Openstack Storage"]
+            },  # Only Openstack Storage SU type is eligible for credit, not HPC
+        )
+        processor.process()
+        output_invoice = processor.data
+
+        expected_invoice = self._get_test_invoice(
+            pi=["PI", "PI"],
+            costs=[50, 75],
+            su_type=["Openstack Storage", "HPC"],
+            credit=[50, None],
+            credit_code=["0005", None],
+            pi_balance=[0, 75],
+            balance=[0, 75],
+        )
+
+        expected_invoice = expected_invoice.astype(output_invoice.dtypes)
+        assert expected_invoice.equals(output_invoice)
+
+    def test_all_eligible_projects(self):
+        """PI has multiple projects and all are eligible, so all should get the credit"""
+        invoice_data = self._get_test_invoice(
+            pi=["PI", "PI"],
+            costs=[40, 60],
+            su_type=["Openstack Storage", "Openstack Compute"],
+        )
+
+        processor = test_utils.new_pi_su_credit_processor(
+            invoice_month="2024-06",
+            data=invoice_data,
+            name="test",
+            pi_su_mapping={
+                "PI": ["Openstack Storage", "Openstack Compute"]
+            },  # Both SU types are eligible for credit
+        )
+        processor.process()
+        output_invoice = processor.data
+
+        expected_invoice = self._get_test_invoice(
+            pi=["PI", "PI"],
+            costs=[40, 60],
+            su_type=["Openstack Storage", "Openstack Compute"],
+            credit=[40, 60],
+            credit_code=["0005", "0005"],
+            pi_balance=[0, 0],
+            balance=[0, 0],
+        )
+
+        expected_invoice = expected_invoice.astype(output_invoice.dtypes)
+        assert expected_invoice.equals(output_invoice)

process_report/tests/util.py

@@ -16,6 +16,7 @@
     bu_subsidy_processor,
     prepayment_processor,
     validate_cluster_name_processor,
+    pi_su_credit_processor,
 )
 
 
@@ -205,3 +206,18 @@ def new_validate_cluster_name_processor(
     return validate_cluster_name_processor.ValidateClusterNameProcessor(
         invoice_month, data, name
     )
+
+
+def new_pi_su_credit_processor(
+    name="",
+    invoice_month="0000-00",
+    data=None,
+    pi_su_mapping=None,
+):
+    if data is None:
+        data = pandas.DataFrame()
+    if pi_su_mapping is None:
+        pi_su_mapping = {}
+    return pi_su_credit_processor.PISUCreditProcessor(
+        invoice_month, data, name, pi_su_mapping
+    )