diff --git a/process_report/invoices/prepay_credits_snapshot.py b/process_report/invoices/prepay_credits_snapshot.py new file mode 100644 index 0000000..8ad1edc --- /dev/null +++ b/process_report/invoices/prepay_credits_snapshot.py @@ -0,0 +1,48 @@ +from dataclasses import dataclass + +import pandas + +from process_report import util +from process_report.invoices import invoice + + +@dataclass +class PrepayCreditsSnapshot(invoice.Invoice): + prepay_credits: pandas.DataFrame + prepay_contacts: pandas.DataFrame + + export_columns_list = [ + invoice.PREPAY_MONTH_FIELD, + invoice.PREPAY_GROUP_NAME_FIELD, + invoice.PREPAY_CREDIT_FIELD, + ] + + @property + def output_path(self): + return f"NERC_Prepaid_Group-Credits-{self.invoice_month}.csv" + + @property + def output_s3_key(self): + return f"Invoices/{self.invoice_month}/NERC_Prepaid_Group-Credits-{self.invoice_month}.csv" + + @property + def output_s3_archive_key(self): + return f"Invoices/{self.invoice_month}/Archive/NERC_Prepaid_Group-Credits-{self.invoice_month} {util.get_iso8601_time()}.csv" + + def _get_prepay_credits_snapshot(self): + managed_groups_list = self.prepay_contacts.loc[ + self.prepay_contacts[invoice.PREPAY_MANAGED_FIELD] == "Yes", + invoice.PREPAY_GROUP_NAME_FIELD, + ] + + credits_mask = ( + self.prepay_credits[invoice.PREPAY_MONTH_FIELD] == self.invoice_month + ) & ( + self.prepay_credits[invoice.PREPAY_GROUP_NAME_FIELD].isin( + managed_groups_list + ) + ) + return self.prepay_credits[credits_mask] + + def _prepare(self): + self.export_data = self._get_prepay_credits_snapshot() diff --git a/process_report/process_report.py b/process_report/process_report.py index 1610a20..38b5585 100644 --- a/process_report/process_report.py +++ b/process_report/process_report.py @@ -16,6 +16,7 @@ bu_internal_invoice, pi_specific_invoice, MOCA_prepaid_invoice, + prepay_credits_snapshot, ) from process_report.processors import ( validate_pi_alias_processor, @@ -358,6 +359,14 @@ def main(): name="", invoice_month=invoice_month, data=processed_data.copy() ) + prepay_credits_snap = prepay_credits_snapshot.PrepayCreditsSnapshot( + name="", + invoice_month=invoice_month, + data=None, + prepay_credits=prepay_credits, + prepay_contacts=prepay_info, + ) + util.process_and_export_invoices( [ lenovo_inv, @@ -367,6 +376,7 @@ def main(): bu_internal_inv, pi_inv, moca_prepaid_inv, + prepay_credits_snap, ], args.upload_to_s3, ) diff --git a/process_report/tests/unit/invoices/test_credits_snapshot.py b/process_report/tests/unit/invoices/test_credits_snapshot.py new file mode 100644 index 0000000..212ea0d --- /dev/null +++ b/process_report/tests/unit/invoices/test_credits_snapshot.py @@ -0,0 +1,41 @@ +from unittest import TestCase +import pandas + +from process_report.tests import util as test_utils + + +class TestCreditsSnapshot(TestCase): + def _get_test_prepay_credits(self, months, group_names, credits): + return pandas.DataFrame( + {"Month": months, "Group Name": group_names, "Credit": credits} + ) + + def _get_test_prepay_contacts(self, group_names, emails, is_managed): + return pandas.DataFrame( + { + "Group Name": group_names, + "Group Contact Email": emails, + "MGHPCC Managed": is_managed, + } + ) + + def test_get_credit_snapshot(self): + invoice_month = "2024-10" + test_prepay_credits = self._get_test_prepay_credits( + ["2024-10", "2024-10", "2024-10", "2024-09", "2024-09"], + ["G1", "G2", "G3", "G1", "G2"], + [0] * 5, + ) + test_prepay_contacts = self._get_test_prepay_contacts( + ["G1", "G2", "G3"], [""] * 3, ["Yes", "No", "Yes"] + ) + answer_credits_snapshot = test_prepay_credits.iloc[[0, 2]] + + new_prepayment_proc = test_utils.new_prepay_credits_snapshot( + invoice_month=invoice_month, + prepay_credits=test_prepay_credits, + prepay_contacts=test_prepay_contacts, + ) + output_snapshot = new_prepayment_proc._get_prepay_credits_snapshot() + + self.assertTrue(answer_credits_snapshot.equals(output_snapshot)) diff --git a/process_report/tests/util.py b/process_report/tests/util.py index 8a39b40..abbe5db 100644 --- a/process_report/tests/util.py +++ b/process_report/tests/util.py @@ -4,6 +4,7 @@ invoice, billable_invoice, pi_specific_invoice, + prepay_credits_snapshot, ) from process_report.processors import ( @@ -175,3 +176,15 @@ def new_prepayment_processor( prepay_debits_filepath, upload_to_s3, ) + + +def new_prepay_credits_snapshot( + name="", + invoice_month="0000-00", + data=None, + prepay_credits=None, + prepay_contacts=None, +): + return prepay_credits_snapshot.PrepayCreditsSnapshot( + name, invoice_month, data, prepay_credits, prepay_contacts + )