From 28fb94338e075f330c84113cf173c4b67b220701 Mon Sep 17 00:00:00 2001 From: Quan Pham Date: Wed, 18 Sep 2024 17:47:03 -0400 Subject: [PATCH 1/5] Initialized processor for applying the New-PI Credit and discounts Code from `BillableInvoice` and `DiscountInvoice` has been copied over without change --- .../processors/discount_processor.py | 78 ++++++++ .../processors/new_pi_credit_processor.py | 174 ++++++++++++++++++ 2 files changed, 252 insertions(+) create mode 100644 process_report/processors/discount_processor.py create mode 100644 process_report/processors/new_pi_credit_processor.py diff --git a/process_report/processors/discount_processor.py b/process_report/processors/discount_processor.py new file mode 100644 index 0000000..8feb24e --- /dev/null +++ b/process_report/processors/discount_processor.py @@ -0,0 +1,78 @@ +from dataclasses import dataclass +import pandas + +from process_report.processors import processor + + +@dataclass +class DiscountProcessor(processor.Processor): + """ + Invoice class containing functions useful for applying discounts + on dataframes + """ + + @staticmethod + def apply_flat_discount( + invoice: pandas.DataFrame, + pi_projects: pandas.DataFrame, + discount_amount: int, + discount_field: str, + balance_field: str, + code_field: str = None, + discount_code: str = None, + ): + """ + Takes in an invoice and a list of PI projects that are a subset of it, + and applies a flat discount to those PI projects. Note that this function + will change the provided `invoice` Dataframe directly. Therefore, it does + not return the changed invoice. + + This function assumes that the balance field shows the remaining cost of the project, + or what the PI would pay before the flat discount is applied. + + If the optional parameters `code_field` and `discount_code` are passed in, + `discount_code` will be comma-APPENDED to the `code_field` of projects where + the discount is applied + + Returns the amount of discount used. + + :param invoice: Dataframe containing all projects + :param pi_projects: A subset of `invoice`, containing all projects for a PI you want to apply the discount + :param discount_amount: The discount given to the PI + :param discount_field: Name of the field to put the discount amount applied to each project + :param balance_field: Name of the balance field + :param code_field: Name of the discount code field + :param discount_code: Code of the discount + """ + + def apply_discount_on_project(remaining_discount_amount, project_i, project): + remaining_project_balance = project[balance_field] + applied_discount = min(remaining_project_balance, remaining_discount_amount) + invoice.at[project_i, discount_field] = applied_discount + invoice.at[project_i, balance_field] = ( + project[balance_field] - applied_discount + ) + remaining_discount_amount -= applied_discount + return remaining_discount_amount + + def apply_credit_code_on_project(project_i): + if code_field and discount_code: + if pandas.isna(invoice.at[project_i, code_field]): + invoice.at[project_i, code_field] = discount_code + else: + invoice.at[project_i, code_field] = ( + invoice.at[project_i, code_field] + "," + discount_code + ) + + remaining_discount_amount = discount_amount + for i, row in pi_projects.iterrows(): + if remaining_discount_amount == 0: + break + else: + remaining_discount_amount = apply_discount_on_project( + remaining_discount_amount, i, row + ) + apply_credit_code_on_project(i) + + discount_used = discount_amount - remaining_discount_amount + return discount_used diff --git a/process_report/processors/new_pi_credit_processor.py b/process_report/processors/new_pi_credit_processor.py new file mode 100644 index 0000000..9b61aa1 --- /dev/null +++ b/process_report/processors/new_pi_credit_processor.py @@ -0,0 +1,174 @@ +import sys + +from dataclasses import dataclass +import pandas +import pyarrow + +from process_report import util +from process_report.invoices import invoice +from process_report.processors import discount_processor + + +@dataclass +class NewPICreditProcessor(discount_processor.DiscountProcessor): + NEW_PI_CREDIT_CODE = "0002" + INITIAL_CREDIT_AMOUNT = 1000 + EXCLUDE_SU_TYPES = ["OpenShift GPUA100SXM4", "OpenStack GPUA100SXM4"] + + old_pi_filepath: str + limit_new_pi_credit_to_partners: bool = False + + @staticmethod + def _load_old_pis(old_pi_filepath) -> pandas.DataFrame: + try: + old_pi_df = pandas.read_csv( + old_pi_filepath, + dtype={ + invoice.PI_INITIAL_CREDITS: pandas.ArrowDtype( + pyarrow.decimal128(21, 2) + ), + invoice.PI_1ST_USED: pandas.ArrowDtype(pyarrow.decimal128(21, 2)), + invoice.PI_2ND_USED: pandas.ArrowDtype(pyarrow.decimal128(21, 2)), + }, + ) + except FileNotFoundError: + sys.exit("Applying credit 0002 failed. Old PI file does not exist") + + return old_pi_df + + @staticmethod + def _get_pi_age(old_pi_df: pandas.DataFrame, pi, invoice_month): + """Returns time difference between current invoice month and PI's first invoice month + I.e 0 for new PIs + Will raise an error if the PI'a age is negative, which suggests a faulty invoice, or a program bug""" + first_invoice_month = old_pi_df.loc[ + old_pi_df[invoice.PI_PI_FIELD] == pi, invoice.PI_FIRST_MONTH + ] + if first_invoice_month.empty: + return 0 + + month_diff = util.get_month_diff(invoice_month, first_invoice_month.iat[0]) + if month_diff < 0: + sys.exit( + f"PI {pi} from {first_invoice_month} found in {invoice_month} invoice!" + ) + else: + return month_diff + + def _filter_partners(self, data): + active_partnerships = list() + institute_list = util.load_institute_list() + for institute_info in institute_list: + if partnership_start_date := institute_info.get( + "mghpcc_partnership_start_date" + ): + if util.get_month_diff(self.invoice_month, partnership_start_date) >= 0: + active_partnerships.append(institute_info["display_name"]) + + return data[data[invoice.INSTITUTION_FIELD].isin(active_partnerships)] + + def _filter_excluded_su_types(self, data): + return data[~(data[invoice.SU_TYPE_FIELD].isin(self.EXCLUDE_SU_TYPES))] + + def _get_credit_eligible_projects(self, data: pandas.DataFrame): + filtered_data = self._filter_excluded_su_types(data) + if self.limit_new_pi_credit_to_partners: + filtered_data = self._filter_partners(filtered_data) + + return filtered_data + + def _apply_credits_new_pi( + self, data: pandas.DataFrame, old_pi_df: pandas.DataFrame + ): + def get_initial_credit_amount( + old_pi_df, invoice_month, default_initial_credit_amount + ): + first_month_processed_pis = old_pi_df[ + old_pi_df[invoice.PI_FIRST_MONTH] == invoice_month + ] + if first_month_processed_pis[ + invoice.PI_INITIAL_CREDITS + ].empty or pandas.isna( + new_pi_credit_amount := first_month_processed_pis[ + invoice.PI_INITIAL_CREDITS + ].iat[0] + ): + new_pi_credit_amount = default_initial_credit_amount + + return new_pi_credit_amount + + new_pi_credit_amount = get_initial_credit_amount( + old_pi_df, self.invoice_month, self.INITIAL_CREDIT_AMOUNT + ) + print(f"New PI Credit set at {new_pi_credit_amount} for {self.invoice_month}") + + credit_eligible_projects = self._get_credit_eligible_projects(data) + current_pi_set = set(credit_eligible_projects[invoice.PI_FIELD]) + for pi in current_pi_set: + pi_projects = credit_eligible_projects[ + credit_eligible_projects[invoice.PI_FIELD] == pi + ] + pi_age = self._get_pi_age(old_pi_df, pi, self.invoice_month) + pi_old_pi_entry = old_pi_df.loc[ + old_pi_df[invoice.PI_PI_FIELD] == pi + ].squeeze() + + if pi_age > 1: + for i, row in pi_projects.iterrows(): + data.at[i, invoice.BALANCE_FIELD] = row[invoice.COST_FIELD] + else: + if pi_age == 0: + if len(pi_old_pi_entry) == 0: + pi_entry = [pi, self.invoice_month, new_pi_credit_amount, 0, 0] + old_pi_df = pandas.concat( + [ + pandas.DataFrame([pi_entry], columns=old_pi_df.columns), + old_pi_df, + ], + ignore_index=True, + ) + pi_old_pi_entry = old_pi_df.loc[ + old_pi_df[invoice.PI_PI_FIELD] == pi + ].squeeze() + + remaining_credit = new_pi_credit_amount + credit_used_field = invoice.PI_1ST_USED + elif pi_age == 1: + remaining_credit = ( + pi_old_pi_entry[invoice.PI_INITIAL_CREDITS] + - pi_old_pi_entry[invoice.PI_1ST_USED] + ) + credit_used_field = invoice.PI_2ND_USED + + credits_used = self.apply_flat_discount( + data, + pi_projects, + remaining_credit, + invoice.CREDIT_FIELD, + invoice.BALANCE_FIELD, + invoice.CREDIT_CODE_FIELD, + self.NEW_PI_CREDIT_CODE, + ) + + if (pi_old_pi_entry[credit_used_field] != 0) and ( + credits_used != pi_old_pi_entry[credit_used_field] + ): + print( + f"Warning: PI file overwritten. PI {pi} previously used ${pi_old_pi_entry[credit_used_field]} of New PI credits, now uses ${credits_used}" + ) + old_pi_df.loc[ + old_pi_df[invoice.PI_PI_FIELD] == pi, credit_used_field + ] = credits_used + + return (data, old_pi_df) + + def _prepare(self): + self.data[invoice.CREDIT_FIELD] = None + self.data[invoice.CREDIT_CODE_FIELD] = None + self.data[invoice.BALANCE_FIELD] = self.data[invoice.COST_FIELD] + self.old_pi_df = self._load_old_pis(self.old_pi_filepath) + + def _process(self): + self.data, self.updated_old_pi_df = self._apply_credits_new_pi( + self.data, self.old_pi_df + ) From 1686888e66b58487c99d4e9edae4fb10c9264700 Mon Sep 17 00:00:00 2001 From: Quan Pham Date: Wed, 18 Sep 2024 12:19:53 -0400 Subject: [PATCH 2/5] Implemented processor for applying the New-PI Credit The `DiscountProcessor` class has been created, to be subclassed by processors which applies discounts, such as the New-PI Credit. This processor class introduces an important class constant, `IS_DISCOUNT_BY_NERC`, which detemines whether the final balance, as opposed to the PI balance (more info below), reflects the discount being applied. The New-PI credit is now implemented by `NewPICreditProcessor` During discussions about billing, it was made noted that some discounts are not provided by the MGHPCC, but instead from other sources, such as the BU subsidy, which is provided to BU PIs by BU. This provided motivation for a `PI Balance` field, which would reflect how much money the PI should be billed, as opposed to the `Balance` field, which currently reflects how much money the MGHPCC should receive. These two fields would not equal each other if the PI received discounts not provided by the MGHPCC. Implementation of `NewPICreditProcessor` and the new billing feature required a range of changes: - `apply_discount_on_project()` in `DiscountProcessor` has been slightly modified, where the PI balance and MGHPCC balance is now calculated seperately. - As `BillableInvoice` no longer performs any processing itself, the dataframe from `NewPICreditProcessor` is now passed to all invoice objects. - The test cases for the New-PI credit have been refactored. Test cases for the new billing feature is not written yet. I plan to write them when the processor for the BU Subsidy is implemented - With the new processor, certain Processor and Invoice classes depend on fields created by other Processors, such as the case with `NewPICreditProcessor` and `ValidateBillablePIsProcessor`. As such, docstrings have been added to indicate dependancies --- process_report/invoices/NERC_total_invoice.py | 9 + process_report/invoices/billable_invoice.py | 179 +---- .../invoices/bu_internal_invoice.py | 9 + process_report/invoices/invoice.py | 1 + .../invoices/pi_specific_invoice.py | 9 + process_report/process_report.py | 40 +- .../processors/discount_processor.py | 22 +- .../processors/new_pi_credit_processor.py | 18 +- process_report/tests/unit_tests.py | 744 +++++++++++------- process_report/tests/util.py | 19 +- 10 files changed, 576 insertions(+), 474 deletions(-) diff --git a/process_report/invoices/NERC_total_invoice.py b/process_report/invoices/NERC_total_invoice.py index 335c52a..92982da 100644 --- a/process_report/invoices/NERC_total_invoice.py +++ b/process_report/invoices/NERC_total_invoice.py @@ -6,6 +6,12 @@ @dataclass class NERCTotalInvoice(invoice.Invoice): + """ + This invoice operates on data processed by these Processors: + - ValidateBillablePIsProcessor + - NewPICreditProcessor + """ + INCLUDED_INSTITUTIONS = [ "Harvard University", "Boston University", @@ -45,6 +51,9 @@ def output_s3_archive_key(self): return f"Invoices/{self.invoice_month}/Archive/NERC-{self.invoice_month}-Total-Invoice {util.get_iso8601_time()}.csv" def _prepare_export(self): + self.data = self.data[ + self.data[invoice.IS_BILLABLE_FIELD] & ~self.data[invoice.MISSING_PI_FIELD] + ] self.data = self.data[ self.data[invoice.INSTITUTION_FIELD].isin(self.INCLUDED_INSTITUTIONS) ].copy() diff --git a/process_report/invoices/billable_invoice.py b/process_report/invoices/billable_invoice.py index bc6408e..1761cc8 100644 --- a/process_report/invoices/billable_invoice.py +++ b/process_report/invoices/billable_invoice.py @@ -1,25 +1,28 @@ from dataclasses import dataclass import logging -import sys import pandas import pyarrow -from process_report.invoices import invoice, discount_invoice -from process_report import util - +from process_report.invoices import invoice logger = logging.getLogger(__name__) logging.basicConfig(level=logging.INFO) @dataclass -class BillableInvoice(discount_invoice.DiscountInvoice): - NEW_PI_CREDIT_CODE = "0002" - INITIAL_CREDIT_AMOUNT = 1000 - EXCLUDE_SU_TYPES = ["OpenShift GPUA100SXM4", "OpenStack GPUA100SXM4"] +class BillableInvoice(invoice.Invoice): + """ + This invoice operates on data processed by these Processors: + - ValidateBillablePIsProcessor + - NewPICreditProcessor + """ + PI_S3_FILEPATH = "PIs/PI.csv" + old_pi_filepath: str + updated_old_pi_df: pandas.DataFrame + export_columns_list = [ invoice.INVOICE_DATE_FIELD, invoice.PROJECT_FIELD, @@ -38,61 +41,10 @@ class BillableInvoice(discount_invoice.DiscountInvoice): invoice.BALANCE_FIELD, ] - old_pi_filepath: str - limit_new_pi_credit_to_partners: bool = False - - @staticmethod - def _load_old_pis(old_pi_filepath) -> pandas.DataFrame: - try: - old_pi_df = pandas.read_csv( - old_pi_filepath, - dtype={ - invoice.PI_INITIAL_CREDITS: pandas.ArrowDtype( - pyarrow.decimal128(21, 2) - ), - invoice.PI_1ST_USED: pandas.ArrowDtype(pyarrow.decimal128(21, 2)), - invoice.PI_2ND_USED: pandas.ArrowDtype(pyarrow.decimal128(21, 2)), - }, - ) - except FileNotFoundError: - sys.exit("Applying credit 0002 failed. Old PI file does not exist") - - return old_pi_df - - @staticmethod - def _get_pi_age(old_pi_df: pandas.DataFrame, pi, invoice_month): - """Returns time difference between current invoice month and PI's first invoice month - I.e 0 for new PIs - Will raise an error if the PI'a age is negative, which suggests a faulty invoice, or a program bug""" - first_invoice_month = old_pi_df.loc[ - old_pi_df[invoice.PI_PI_FIELD] == pi, invoice.PI_FIRST_MONTH - ] - if first_invoice_month.empty: - return 0 - - month_diff = util.get_month_diff(invoice_month, first_invoice_month.iat[0]) - if month_diff < 0: - sys.exit( - f"PI {pi} from {first_invoice_month} found in {invoice_month} invoice!" - ) - else: - return month_diff - - def _prepare(self): + def _prepare_export(self): self.data = self.data[ self.data[invoice.IS_BILLABLE_FIELD] & ~self.data[invoice.MISSING_PI_FIELD] ] - self.data[invoice.CREDIT_FIELD] = None - self.data[invoice.CREDIT_CODE_FIELD] = None - self.data[invoice.BALANCE_FIELD] = self.data[invoice.COST_FIELD] - self.old_pi_df = self._load_old_pis(self.old_pi_filepath) - - def _process(self): - self.data, self.updated_old_pi_df = self._apply_credits_new_pi( - self.data, self.old_pi_df - ) - - def _prepare_export(self): self.updated_old_pi_df = self.updated_old_pi_df.astype( { invoice.PI_INITIAL_CREDITS: pandas.ArrowDtype( @@ -110,110 +62,3 @@ def export(self): def export_s3(self, s3_bucket): super().export_s3(s3_bucket) s3_bucket.upload_file(self.old_pi_filepath, self.PI_S3_FILEPATH) - - def _filter_partners(self, data): - active_partnerships = list() - institute_list = util.load_institute_list() - for institute_info in institute_list: - if partnership_start_date := institute_info.get( - "mghpcc_partnership_start_date" - ): - if util.get_month_diff(self.invoice_month, partnership_start_date) >= 0: - active_partnerships.append(institute_info["display_name"]) - - return data[data[invoice.INSTITUTION_FIELD].isin(active_partnerships)] - - def _filter_excluded_su_types(self, data): - return data[~(data[invoice.SU_TYPE_FIELD].isin(self.EXCLUDE_SU_TYPES))] - - def _get_credit_eligible_projects(self, data: pandas.DataFrame): - filtered_data = self._filter_excluded_su_types(data) - if self.limit_new_pi_credit_to_partners: - filtered_data = self._filter_partners(filtered_data) - - return filtered_data - - def _apply_credits_new_pi( - self, data: pandas.DataFrame, old_pi_df: pandas.DataFrame - ): - def get_initial_credit_amount( - old_pi_df, invoice_month, default_initial_credit_amount - ): - first_month_processed_pis = old_pi_df[ - old_pi_df[invoice.PI_FIRST_MONTH] == invoice_month - ] - if first_month_processed_pis[ - invoice.PI_INITIAL_CREDITS - ].empty or pandas.isna( - new_pi_credit_amount := first_month_processed_pis[ - invoice.PI_INITIAL_CREDITS - ].iat[0] - ): - new_pi_credit_amount = default_initial_credit_amount - - return new_pi_credit_amount - - new_pi_credit_amount = get_initial_credit_amount( - old_pi_df, self.invoice_month, self.INITIAL_CREDIT_AMOUNT - ) - print(f"New PI Credit set at {new_pi_credit_amount} for {self.invoice_month}") - - credit_eligible_projects = self._get_credit_eligible_projects(data) - current_pi_set = set(credit_eligible_projects[invoice.PI_FIELD]) - for pi in current_pi_set: - pi_projects = credit_eligible_projects[ - credit_eligible_projects[invoice.PI_FIELD] == pi - ] - pi_age = self._get_pi_age(old_pi_df, pi, self.invoice_month) - pi_old_pi_entry = old_pi_df.loc[ - old_pi_df[invoice.PI_PI_FIELD] == pi - ].squeeze() - - if pi_age > 1: - for i, row in pi_projects.iterrows(): - data.at[i, invoice.BALANCE_FIELD] = row[invoice.COST_FIELD] - else: - if pi_age == 0: - if len(pi_old_pi_entry) == 0: - pi_entry = [pi, self.invoice_month, new_pi_credit_amount, 0, 0] - old_pi_df = pandas.concat( - [ - pandas.DataFrame([pi_entry], columns=old_pi_df.columns), - old_pi_df, - ], - ignore_index=True, - ) - pi_old_pi_entry = old_pi_df.loc[ - old_pi_df[invoice.PI_PI_FIELD] == pi - ].squeeze() - - remaining_credit = new_pi_credit_amount - credit_used_field = invoice.PI_1ST_USED - elif pi_age == 1: - remaining_credit = ( - pi_old_pi_entry[invoice.PI_INITIAL_CREDITS] - - pi_old_pi_entry[invoice.PI_1ST_USED] - ) - credit_used_field = invoice.PI_2ND_USED - - credits_used = self.apply_flat_discount( - data, - pi_projects, - remaining_credit, - invoice.CREDIT_FIELD, - invoice.BALANCE_FIELD, - invoice.CREDIT_CODE_FIELD, - self.NEW_PI_CREDIT_CODE, - ) - - if (pi_old_pi_entry[credit_used_field] != 0) and ( - credits_used != pi_old_pi_entry[credit_used_field] - ): - print( - f"Warning: PI file overwritten. PI {pi} previously used ${pi_old_pi_entry[credit_used_field]} of New PI credits, now uses ${credits_used}" - ) - old_pi_df.loc[ - old_pi_df[invoice.PI_PI_FIELD] == pi, credit_used_field - ] = credits_used - - return (data, old_pi_df) diff --git a/process_report/invoices/bu_internal_invoice.py b/process_report/invoices/bu_internal_invoice.py index bc6f9c2..e1f7732 100644 --- a/process_report/invoices/bu_internal_invoice.py +++ b/process_report/invoices/bu_internal_invoice.py @@ -7,6 +7,12 @@ @dataclass class BUInternalInvoice(discount_invoice.DiscountInvoice): + """ + This invoice operates on data processed by these Processors: + - ValidateBillablePIsProcessor + - NewPICreditProcessor + """ + export_columns_list = [ invoice.INVOICE_DATE_FIELD, invoice.PI_FIELD, @@ -27,6 +33,9 @@ def get_project(row): else: return project_alloc[: project_alloc.rfind("-")] + self.data = self.data[ + self.data[invoice.IS_BILLABLE_FIELD] & ~self.data[invoice.MISSING_PI_FIELD] + ] self.data = self.data[ self.data[invoice.INSTITUTION_FIELD] == "Boston University" ].copy() diff --git a/process_report/invoices/invoice.py b/process_report/invoices/invoice.py index 3b5b2e5..35c8ffb 100644 --- a/process_report/invoices/invoice.py +++ b/process_report/invoices/invoice.py @@ -36,6 +36,7 @@ ### Internally used field names IS_BILLABLE_FIELD = "Is Billable" MISSING_PI_FIELD = "Missing PI" +PI_BALANCE_FIELD = "PI Balance" ### diff --git a/process_report/invoices/pi_specific_invoice.py b/process_report/invoices/pi_specific_invoice.py index 6d52933..dabc0a0 100644 --- a/process_report/invoices/pi_specific_invoice.py +++ b/process_report/invoices/pi_specific_invoice.py @@ -9,6 +9,12 @@ @dataclass class PIInvoice(invoice.Invoice): + """ + This invoice operates on data processed by these Processors: + - ValidateBillablePIsProcessor + - NewPICreditProcessor + """ + export_columns_list = [ invoice.INVOICE_DATE_FIELD, invoice.PROJECT_FIELD, @@ -28,6 +34,9 @@ class PIInvoice(invoice.Invoice): ] def _prepare(self): + self.data = self.data[ + self.data[invoice.IS_BILLABLE_FIELD] & ~self.data[invoice.MISSING_PI_FIELD] + ] self.pi_list = self.data[invoice.PI_FIELD].unique() def export(self): diff --git a/process_report/process_report.py b/process_report/process_report.py index 6ee32f1..cd65430 100644 --- a/process_report/process_report.py +++ b/process_report/process_report.py @@ -20,6 +20,7 @@ add_institution_processor, lenovo_processor, validate_billable_pi_processor, + new_pi_credit_processor, ) ### PI file field names @@ -229,7 +230,19 @@ def main(): ) validate_billable_pi_proc.process() - processed_data = validate_billable_pi_proc.data + rates_info = load_from_url() + new_pi_credit_proc = new_pi_credit_processor.NewPICreditProcessor( + "", + invoice_month, + data=validate_billable_pi_proc.data, + old_pi_filepath=old_pi_file, + limit_new_pi_credit_to_partners=rates_info.get_value_at( + "Limit New PI Credit to MGHPCC Partners", invoice_month + ), + ) + new_pi_credit_proc.process() + + processed_data = new_pi_credit_proc.data ### Initialize invoices @@ -249,40 +262,41 @@ def main(): if args.upload_to_s3: backup_to_s3_old_pi_file(old_pi_file) - rates_info = load_from_url() billable_inv = billable_invoice.BillableInvoice( name=args.output_file, invoice_month=invoice_month, data=processed_data.copy(), old_pi_filepath=old_pi_file, - limit_new_pi_credit_to_partners=rates_info.get_value_at( - "Limit New PI Credit to MGHPCC Partners", invoice_month - ), - ) - - util.process_and_export_invoices( - [lenovo_inv, nonbillable_inv, billable_inv], args.upload_to_s3 + updated_old_pi_df=new_pi_credit_proc.updated_old_pi_df, ) nerc_total_inv = NERC_total_invoice.NERCTotalInvoice( name=args.NERC_total_invoice_file, invoice_month=invoice_month, - data=billable_inv.data.copy(), + data=processed_data.copy(), ) bu_internal_inv = bu_internal_invoice.BUInternalInvoice( name=args.BU_invoice_file, invoice_month=invoice_month, - data=billable_inv.data.copy(), + data=processed_data.copy(), subsidy_amount=args.BU_subsidy_amount, ) pi_inv = pi_specific_invoice.PIInvoice( - name=args.output_folder, invoice_month=invoice_month, data=billable_inv.data + name=args.output_folder, invoice_month=invoice_month, data=processed_data.copy() ) util.process_and_export_invoices( - [nerc_total_inv, bu_internal_inv, pi_inv], args.upload_to_s3 + [ + lenovo_inv, + nonbillable_inv, + billable_inv, + nerc_total_inv, + bu_internal_inv, + pi_inv, + ], + args.upload_to_s3, ) diff --git a/process_report/processors/discount_processor.py b/process_report/processors/discount_processor.py index 8feb24e..ec089e8 100644 --- a/process_report/processors/discount_processor.py +++ b/process_report/processors/discount_processor.py @@ -1,21 +1,22 @@ -from dataclasses import dataclass import pandas from process_report.processors import processor -@dataclass class DiscountProcessor(processor.Processor): """ - Invoice class containing functions useful for applying discounts + Processor class containing functions useful for applying discounts on dataframes """ - @staticmethod + IS_DISCOUNT_BY_NERC = True + def apply_flat_discount( + self, invoice: pandas.DataFrame, pi_projects: pandas.DataFrame, - discount_amount: int, + pi_balance_field: str, + discount_amount: float, discount_field: str, balance_field: str, code_field: str = None, @@ -38,20 +39,21 @@ def apply_flat_discount( :param invoice: Dataframe containing all projects :param pi_projects: A subset of `invoice`, containing all projects for a PI you want to apply the discount + :param pi_balance_field: Name of the field of the PI balance :param discount_amount: The discount given to the PI :param discount_field: Name of the field to put the discount amount applied to each project - :param balance_field: Name of the balance field + :param balance_field: Name of the NERC balance field :param code_field: Name of the discount code field :param discount_code: Code of the discount """ def apply_discount_on_project(remaining_discount_amount, project_i, project): - remaining_project_balance = project[balance_field] + remaining_project_balance = project[pi_balance_field] applied_discount = min(remaining_project_balance, remaining_discount_amount) invoice.at[project_i, discount_field] = applied_discount - invoice.at[project_i, balance_field] = ( - project[balance_field] - applied_discount - ) + invoice.at[project_i, pi_balance_field] -= applied_discount + if self.IS_DISCOUNT_BY_NERC: + invoice.at[project_i, balance_field] -= applied_discount remaining_discount_amount -= applied_discount return remaining_discount_amount diff --git a/process_report/processors/new_pi_credit_processor.py b/process_report/processors/new_pi_credit_processor.py index 9b61aa1..576c50b 100644 --- a/process_report/processors/new_pi_credit_processor.py +++ b/process_report/processors/new_pi_credit_processor.py @@ -11,9 +11,15 @@ @dataclass class NewPICreditProcessor(discount_processor.DiscountProcessor): + """ + This processor operates on data processed by these Processors: + - ValidateBillablePIsProcessor + """ + NEW_PI_CREDIT_CODE = "0002" INITIAL_CREDIT_AMOUNT = 1000 EXCLUDE_SU_TYPES = ["OpenShift GPUA100SXM4", "OpenStack GPUA100SXM4"] + IS_DISCOUNT_BY_NERC = True old_pi_filepath: str limit_new_pi_credit_to_partners: bool = False @@ -70,8 +76,16 @@ def _filter_partners(self, data): def _filter_excluded_su_types(self, data): return data[~(data[invoice.SU_TYPE_FIELD].isin(self.EXCLUDE_SU_TYPES))] + def _filter_nonbillables(self, data): + return data[data["Is Billable"]] + + def _filter_missing_pis(self, data): + return data[~data["Missing PI"]] + def _get_credit_eligible_projects(self, data: pandas.DataFrame): - filtered_data = self._filter_excluded_su_types(data) + filtered_data = self._filter_nonbillables(data) + filtered_data = self._filter_missing_pis(filtered_data) + filtered_data = self._filter_excluded_su_types(filtered_data) if self.limit_new_pi_credit_to_partners: filtered_data = self._filter_partners(filtered_data) @@ -143,6 +157,7 @@ def get_initial_credit_amount( credits_used = self.apply_flat_discount( data, pi_projects, + invoice.PI_BALANCE_FIELD, remaining_credit, invoice.CREDIT_FIELD, invoice.BALANCE_FIELD, @@ -165,6 +180,7 @@ def get_initial_credit_amount( def _prepare(self): self.data[invoice.CREDIT_FIELD] = None self.data[invoice.CREDIT_CODE_FIELD] = None + self.data[invoice.PI_BALANCE_FIELD] = self.data[invoice.COST_FIELD] self.data[invoice.BALANCE_FIELD] = self.data[invoice.COST_FIELD] self.old_pi_df = self._load_old_pis(self.old_pi_filepath) diff --git a/process_report/tests/unit_tests.py b/process_report/tests/unit_tests.py index 047a325..f09a7de 100644 --- a/process_report/tests/unit_tests.py +++ b/process_report/tests/unit_tests.py @@ -1,7 +1,6 @@ from unittest import TestCase, mock import tempfile import pandas -import pyarrow import os import uuid import math @@ -110,6 +109,8 @@ def setUp(self): "ProjectE", ], "Untouch Data Column": ["DataA", "DataB", "DataC", "DataD", "DataE"], + "Is Billable": [True, True, True, True, True], + "Missing PI": [False, False, False, False, False], } self.dataframe = pandas.DataFrame(data) self.invoice_month = data["Invoice Month"][0] @@ -275,304 +276,483 @@ def test_get_month_diff(self): util.get_month_diff("2024-16", "2025-03") -class TestCredit0002(TestCase): - def setUp(self): - data = { - "Invoice Month": [ - "2024-03", - "2024-03", - "2024-03", - "2024-03", - "2024-03", - "2024-03", - "2024-03", - "2024-03", - "2024-03", - "2024-03", - "2024-03", - ], - "Manager (PI)": [ - "PI1", - "PI2", - "PI3", - "PI4", - "PI4", - "PI5", - "PI7", - "NewPI1", - "NewPI1", - "NewPI2", - "NewPI2", - ], - "SU Type": [ - "CPU", - "CPU", - "CPU", - "CPU", - "CPU", - "CPU", - "CPU", - "CPU", - "CPU", - "CPU", - "CPU", +class TestNewPICreditProcessor(TestCase): + def _assert_result_invoice_and_old_pi_file( + self, + invoice_month, + test_invoice, + test_old_pi_filepath, + answer_invoice, + answer_old_pi_df, + ): + new_pi_credit_proc = test_utils.new_new_pi_credit_processor( + invoice_month=invoice_month, + data=test_invoice, + old_pi_filepath=test_old_pi_filepath, + ) + new_pi_credit_proc.process() + output_invoice = new_pi_credit_proc.data + output_old_pi_df = new_pi_credit_proc.updated_old_pi_df.sort_values( + by="PI", ignore_index=True + ) + + answer_invoice = answer_invoice.astype(output_invoice.dtypes) + answer_old_pi_df = answer_old_pi_df.astype(output_old_pi_df.dtypes).sort_values( + by="PI", ignore_index=True + ) + + self.assertTrue(output_invoice.equals(answer_invoice)) + self.assertTrue(output_old_pi_df.equals(answer_old_pi_df)) + + def _get_test_invoice( + self, pi, cost, su_type=None, is_billable=None, missing_pi=None + ): + if not su_type: + su_type = ["CPU" for _ in range(len(pi))] + + if not is_billable: + is_billable = [True for _ in range(len(pi))] + + if not missing_pi: + missing_pi = [False for _ in range(len(pi))] + + return pandas.DataFrame( + { + "Manager (PI)": pi, + "Cost": cost, + "SU Type": su_type, + "Is Billable": is_billable, + "Missing PI": missing_pi, + } + ) + + def setUp(self) -> None: + self.test_old_pi_file = tempfile.NamedTemporaryFile( + delete=False, mode="w+", suffix=".csv" + ) + + def tearDown(self) -> None: + os.remove(self.test_old_pi_file.name) + + def test_no_new_pi(self): + test_invoice = self._get_test_invoice( + ["PI" for _ in range(3)], [100 for _ in range(3)] + ) + + # Other fields of old PI file not accessed if PI is no longer + # eligible for new-PI credit + test_old_pi_df = pandas.DataFrame( + { + "PI": ["PI"], + "First Invoice Month": ["2024-01"], + "Initial Credits": [1000], + } + ) + test_old_pi_df.to_csv(self.test_old_pi_file.name, index=False) + + answer_invoice = pandas.concat( + [ + test_invoice, + pandas.DataFrame( + { + "Credit": [None for _ in range(3)], + "Credit Code": [None for _ in range(3)], + "PI Balance": [100 for _ in range(3)], + "Balance": [100 for _ in range(3)], + } + ), ], - "Project - Allocation": [ - "ProjectA", - "ProjectB", - "ProjectC", - "ProjectD", - "ProjectE", - "ProjectF", - "ProjectG", - "ProjectH", - "ProjectI", - "ProjectJ", - "ProjectK", + axis=1, + ) + + answer_old_pi_df = test_old_pi_df.copy() + + self._assert_result_invoice_and_old_pi_file( + "2024-06", + test_invoice, + self.test_old_pi_file.name, + answer_invoice, + answer_old_pi_df, + ) + + def test_one_new_pi(self): + """Invoice with one completely new PI""" + + # One allocation + invoice_month = "2024-06" + + test_invoice = self._get_test_invoice(["PI"], [100]) + + test_old_pi_df = pandas.DataFrame( + columns=[ + "PI", + "First Invoice Month", + "Initial Credits", + "1st Month Used", + "2nd Month Used", + ] + ) + test_old_pi_df.to_csv(self.test_old_pi_file.name, index=False) + + answer_invoice = pandas.concat( + [ + test_invoice, + pandas.DataFrame( + { + "Credit": [100], + "Credit Code": ["0002"], + "PI Balance": [0], + "Balance": [0], + } + ), ], - "Cost": [10, 100, 10000, 500, 100, 400, 200, 250, 250, 700, 700], - } - answer_df_dict = { - "Invoice Month": [ - "2024-03", - "2024-03", - "2024-03", - "2024-03", - "2024-03", - "2024-03", - "2024-03", - "2024-03", - "2024-03", - "2024-03", - "2024-03", + axis=1, + ) + + answer_old_pi_df = pandas.DataFrame( + { + "PI": ["PI"], + "First Invoice Month": ["2024-06"], + "Initial Credits": [1000], + "1st Month Used": [100], + "2nd Month Used": [0], + } + ) + + self._assert_result_invoice_and_old_pi_file( + invoice_month, + test_invoice, + self.test_old_pi_file.name, + answer_invoice, + answer_old_pi_df, + ) + + # Two allocations, costs partially covered + test_invoice = self._get_test_invoice(["PI", "PI"], [500, 1000]) + + answer_invoice = pandas.concat( + [ + test_invoice, + pandas.DataFrame( + { + "Credit": [500, 500], + "Credit Code": ["0002", "0002"], + "PI Balance": [0, 500], + "Balance": [0, 500], + } + ), ], - "Manager (PI)": [ - "PI1", - "PI2", - "PI3", - "PI4", - "PI4", - "PI5", - "PI7", - "NewPI1", - "NewPI1", - "NewPI2", - "NewPI2", + axis=1, + ) + + answer_old_pi_df = pandas.DataFrame( + { + "PI": ["PI"], + "First Invoice Month": ["2024-06"], + "Initial Credits": [1000], + "1st Month Used": [1000], + "2nd Month Used": [0], + } + ) + + self._assert_result_invoice_and_old_pi_file( + invoice_month, + test_invoice, + self.test_old_pi_file.name, + answer_invoice, + answer_old_pi_df, + ) + + # Two allocations, costs completely covered + test_invoice = self._get_test_invoice(["PI", "PI"], [500, 400]) + + answer_invoice = pandas.concat( + [ + test_invoice, + pandas.DataFrame( + { + "Credit": [500, 400], + "Credit Code": ["0002", "0002"], + "PI Balance": [0, 0], + "Balance": [0, 0], + } + ), ], - "SU Type": [ - "CPU", - "CPU", - "CPU", - "CPU", - "CPU", - "CPU", - "CPU", - "CPU", - "CPU", - "CPU", - "CPU", + axis=1, + ) + + answer_old_pi_df = pandas.DataFrame( + { + "PI": ["PI"], + "First Invoice Month": ["2024-06"], + "Initial Credits": [1000], + "1st Month Used": [900], + "2nd Month Used": [0], + } + ) + + self._assert_result_invoice_and_old_pi_file( + invoice_month, + test_invoice, + self.test_old_pi_file.name, + answer_invoice, + answer_old_pi_df, + ) + + def test_one_month_pi(self): + """PI has appeared in invoices for one month""" + + # Remaining credits completely covers costs + invoice_month = "2024-07" + test_invoice = self._get_test_invoice(["PI"], [200]) + + test_old_pi_df = pandas.DataFrame( + { + "PI": ["PI"], + "First Invoice Month": ["2024-06"], + "Initial Credits": [1000], + "1st Month Used": [500], + "2nd Month Used": [0], + } + ) + test_old_pi_df.to_csv(self.test_old_pi_file.name, index=False) + + answer_invoice = pandas.concat( + [ + test_invoice, + pandas.DataFrame( + { + "Credit": [200], + "Credit Code": ["0002"], + "PI Balance": [0], + "Balance": [0], + } + ), ], - "Project - Allocation": [ - "ProjectA", - "ProjectB", - "ProjectC", - "ProjectD", - "ProjectE", - "ProjectF", - "ProjectG", - "ProjectH", - "ProjectI", - "ProjectJ", - "ProjectK", + axis=1, + ) + + answer_old_pi_df = pandas.DataFrame( + { + "PI": ["PI"], + "First Invoice Month": ["2024-06"], + "Initial Credits": [1000], + "1st Month Used": [500], + "2nd Month Used": [200], + } + ) + + self._assert_result_invoice_and_old_pi_file( + invoice_month, + test_invoice, + self.test_old_pi_file.name, + answer_invoice, + answer_old_pi_df, + ) + + # Remaining credits partially covers costs + test_invoice = self._get_test_invoice(["PI"], [600]) + + answer_invoice = pandas.concat( + [ + test_invoice, + pandas.DataFrame( + { + "Credit": [500], + "Credit Code": ["0002"], + "PI Balance": [100], + "Balance": [100], + } + ), ], - "Cost": [10, 100, 10000, 500, 100, 400, 200, 250, 250, 700, 700], - "Credit": [None, None, None, 100, None, 400, 200, 250, 250, 500, None], - "Credit Code": [ - None, - None, - None, - "0002", - None, - "0002", - "0002", - "0002", - "0002", - "0002", - None, + axis=1, + ) + + answer_old_pi_df = pandas.DataFrame( + { + "PI": ["PI"], + "First Invoice Month": ["2024-06"], + "Initial Credits": [1000], + "1st Month Used": [500], + "2nd Month Used": [500], + } + ) + + self._assert_result_invoice_and_old_pi_file( + invoice_month, + test_invoice, + self.test_old_pi_file.name, + answer_invoice, + answer_old_pi_df, + ) + + def test_two_new_pi(self): + """Two PIs of different age""" + + # Costs partially and completely covered + invoice_month = "2024-07" + test_invoice = self._get_test_invoice(["PI1", "PI1", "PI2"], [800, 500, 500]) + + test_old_pi_df = pandas.DataFrame( + { + "PI": ["PI1"], + "First Invoice Month": ["2024-06"], + "Initial Credits": [1000], + "1st Month Used": [500], + "2nd Month Used": [0], + } + ) + test_old_pi_df.to_csv(self.test_old_pi_file.name, index=False) + + answer_invoice = pandas.concat( + [ + test_invoice, + pandas.DataFrame( + { + "Credit": [500, None, 500], + "Credit Code": ["0002", None, "0002"], + "PI Balance": [300, 500, 0], + "Balance": [300, 500, 0], + } + ), ], - "Balance": [10, 100, 10000, 400, 100, 0, 0, 0, 0, 200, 700], - } - self.dataframe = pandas.DataFrame(data) - self.dataframe["Credit"] = None - self.dataframe["Credit Code"] = None - self.dataframe["Balance"] = self.dataframe["Cost"] - self.answer_dataframe = pandas.DataFrame(answer_df_dict) - old_pi = [ - "PI,First Invoice Month,Initial Credits,1st Month Used,2nd Month Used", - "PI1,2023-09,500,200,0", - "PI2,2024-01,2000,0,0", - "PI3,2024-01,2000,1000,500", - "PI4,2024-02,1000,900,0", - "PI5,2024-02,1000,300,500", - "PI6,2024-02,1000,700,0", - "PI7,2024-03,500,300,0", # This as current month we're testing, new PIs should get $500 - "PI8,2024-04,1000,500,0", - ] - self.old_pi_df_answer = ( - pandas.DataFrame( - { - "PI": [ - "PI1", - "PI2", - "PI3", - "PI4", - "PI5", - "PI6", - "PI7", - "NewPI1", - "NewPI2", - "PI8", - ], - "First Invoice Month": [ - "2023-09", - "2024-01", - "2024-01", - "2024-02", - "2024-02", - "2024-02", - "2024-03", - "2024-03", - "2024-03", - "2024-04", - ], - "Initial Credits": [ - 500, - 2000, - 2000, - 1000, - 1000, - 1000, - 500, - 500, - 500, - 1000, - ], - "1st Month Used": [200, 0, 1000, 900, 300, 700, 200, 500, 500, 500], - "2nd Month Used": [0, 0, 500, 100, 400, 0, 0, 0, 0, 0], - } - ) - .astype( - { - "Initial Credits": pandas.ArrowDtype(pyarrow.decimal128(21, 2)), - "1st Month Used": pandas.ArrowDtype(pyarrow.decimal128(21, 2)), - "2nd Month Used": pandas.ArrowDtype(pyarrow.decimal128(21, 2)), - }, - ) - .sort_values(by="PI", ignore_index=True) + axis=1, ) - # Contains cases with new, one month old, two month old, older PI, and future PI that hasn't appeared in invoices yet - # For each invoice month, test case where pi has 1 project, >1, and has spare credit - old_pi_file = tempfile.NamedTemporaryFile( - delete=False, mode="w+", suffix=".csv" + answer_old_pi_df = pandas.DataFrame( + { + "PI": ["PI1", "PI2"], + "First Invoice Month": ["2024-06", "2024-07"], + "Initial Credits": [1000, 1000], + "1st Month Used": [500, 500], + "2nd Month Used": [500, 0], + } + ) + + self._assert_result_invoice_and_old_pi_file( + invoice_month, + test_invoice, + self.test_old_pi_file.name, + answer_invoice, + answer_old_pi_df, ) - for pi in old_pi: - old_pi_file.write(pi + "\n") - self.old_pi_file = old_pi_file.name - self.dataframe_no_gpu = pandas.DataFrame( + def test_old_pi_file_overwritten(self): + """If PI already has entry in Old PI file, + their initial credits and PI entry could be overwritten""" + + invoice_month = "2024-06" + test_invoice = self._get_test_invoice(["PI", "PI"], [500, 500]) + test_old_pi_df = pandas.DataFrame( { - "Invoice Month": [ - "2024-03", - "2024-03", - "2024-03", - "2024-03", - "2024-03", - ], - "Manager (PI)": ["PI1", "PI1", "PI1", "PI2", "PI2"], - "SU Type": [ - "GPU", - "OpenShift GPUA100SXM4", - "OpenStack GPUA100SXM4", - "OpenShift GPUA100SXM4", - "OpenStack GPUA100SXM4", - ], - "Cost": [500, 100, 100, 500, 500], + "PI": ["PI"], + "First Invoice Month": ["2024-06"], + "Initial Credits": [500], + "1st Month Used": [200], + "2nd Month Used": [0], } ) - self.dataframe_no_gpu["Credit"] = None - self.dataframe_no_gpu["Credit Code"] = None - self.dataframe_no_gpu["Balance"] = self.dataframe_no_gpu["Cost"] - old_pi_no_gpu = [ - "PI,First Invoice Month,Initial Credits,1st Month Used,2nd Month Used", - "OldPI,2024-03,500,200,0", - ] - old_pi_no_gpu_file = tempfile.NamedTemporaryFile( - delete=False, mode="w", suffix=".csv" + test_old_pi_df.to_csv(self.test_old_pi_file.name, index=False) + + answer_invoice = pandas.concat( + [ + test_invoice, + pandas.DataFrame( + { + "Credit": [500, None], + "Credit Code": ["0002", None], + "PI Balance": [0, 500], + "Balance": [0, 500], + } + ), + ], + axis=1, ) - for pi in old_pi_no_gpu: - old_pi_no_gpu_file.write(pi + "\n") - self.old_pi_no_gpu_file = old_pi_no_gpu_file.name - self.no_gpu_df_answer = pandas.DataFrame( + + answer_old_pi_df = pandas.DataFrame( { - "Invoice Month": [ - "2024-03", - "2024-03", - "2024-03", - "2024-03", - "2024-03", - ], - "Manager (PI)": ["PI1", "PI1", "PI1", "PI2", "PI2"], - "SU Type": [ - "GPU", - "OpenShift GPUA100SXM4", - "OpenStack GPUA100SXM4", - "OpenShift GPUA100SXM4", - "OpenStack GPUA100SXM4", - ], - "Cost": [500, 100, 100, 500, 500], - "Credit": [500, None, None, None, None], - "Credit Code": ["0002", None, None, None, None], - "Balance": [0.0, 100.0, 100.0, 500.0, 500.0], + "PI": ["PI"], + "First Invoice Month": ["2024-06"], + "Initial Credits": [500], + "1st Month Used": [500], + "2nd Month Used": [0], } ) - def tearDown(self): - os.remove(self.old_pi_file) - os.remove(self.old_pi_no_gpu_file) - - def test_apply_credit_0002(self): - test_invoice = test_utils.new_billable_invoice(invoice_month="2024-03") - old_pi_df = test_invoice._load_old_pis(self.old_pi_file) - dataframe, updated_old_pi_df = test_invoice._apply_credits_new_pi( - self.dataframe, old_pi_df - ) - dataframe = dataframe.astype({"Credit": "float64", "Balance": "int64"}) - updated_old_pi_df = updated_old_pi_df.astype( - dtype={ - "Initial Credits": pandas.ArrowDtype(pyarrow.decimal128(21, 2)), - "1st Month Used": pandas.ArrowDtype(pyarrow.decimal128(21, 2)), - "2nd Month Used": pandas.ArrowDtype(pyarrow.decimal128(21, 2)), - }, - ).sort_values(by=["PI"], ignore_index=True) - self.assertTrue(self.answer_dataframe.equals(dataframe)) - self.assertTrue(self.old_pi_df_answer.equals(updated_old_pi_df)) - - def test_no_gpu(self): - test_invoice = test_utils.new_billable_invoice(invoice_month="2024-03") - old_pi_df = test_invoice._load_old_pis(self.old_pi_no_gpu_file) - dataframe, _ = test_invoice._apply_credits_new_pi( - self.dataframe_no_gpu, old_pi_df - ) - dataframe = dataframe.astype({"Credit": "float64", "Balance": "float64"}) - self.assertTrue(self.no_gpu_df_answer.equals(dataframe)) + self._assert_result_invoice_and_old_pi_file( + invoice_month, + test_invoice, + self.test_old_pi_file.name, + answer_invoice, + answer_old_pi_df, + ) + + def test_excluded_su_types(self): + """Certain SU types can be excluded from the credit""" + + invoice_month = "2024-06" + test_invoice = self._get_test_invoice( + ["PI", "PI", "PI", "PI"], + [600, 600, 600, 600], + [ + "CPU", + "OpenShift GPUA100SXM4", + "GPU", + "OpenStack GPUA100SXM4", + ], + ) + + test_old_pi_df = pandas.DataFrame( + columns=[ + "PI", + "First Invoice Month", + "Initial Credits", + "1st Month Used", + "2nd Month Used", + ] + ) + test_old_pi_df.to_csv(self.test_old_pi_file.name, index=False) + + answer_invoice = pandas.concat( + [ + test_invoice, + pandas.DataFrame( + { + "Credit": [600, None, 400, None], + "Credit Code": ["0002", None, "0002", None], + "PI Balance": [0, 600, 200, 600], + "Balance": [0, 600, 200, 600], + } + ), + ], + axis=1, + ) + + answer_old_pi_df = pandas.DataFrame( + { + "PI": ["PI"], + "First Invoice Month": ["2024-06"], + "Initial Credits": [1000], + "1st Month Used": [1000], + "2nd Month Used": [0], + } + ) + + self._assert_result_invoice_and_old_pi_file( + invoice_month, + test_invoice, + self.test_old_pi_file.name, + answer_invoice, + answer_old_pi_df, + ) def test_apply_credit_error(self): + """Test faulty data""" old_pi_df = pandas.DataFrame( {"PI": ["PI1"], "First Invoice Month": ["2024-04"]} ) invoice_month = "2024-03" - test_invoice = test_utils.new_billable_invoice() + test_invoice = test_utils.new_new_pi_credit_processor() with self.assertRaises(SystemExit): test_invoice._get_pi_age(old_pi_df, "PI1", invoice_month) @@ -632,6 +812,8 @@ def setUp(self): 50, 100, ], # Test case where subsidy does/doesn't cover fully balance + "Is Billable": [True, True, True, True, True, True, True, True], + "Missing PI": [False, False, False, False, False, False, False, False], } self.dataframe = pandas.DataFrame(data) self.subsidy = 100 @@ -768,24 +950,24 @@ def test_flag_limit_new_pi_credit(self, mock_load_institute_list): "Institution": ["BU", "HU", "NEU", "MIT", "BC"], } ) - sample_inv = test_utils.new_billable_invoice( + sample_proc = test_utils.new_new_pi_credit_processor( limit_new_pi_credit_to_partners=True ) # When no partnerships are active - sample_inv.invoice_month = "2024-01" - output_df = sample_inv._filter_partners(sample_df) + sample_proc.invoice_month = "2024-01" + output_df = sample_proc._filter_partners(sample_df) self.assertTrue(output_df.empty) # When some partnerships are active - sample_inv.invoice_month = "2024-06" - output_df = sample_inv._filter_partners(sample_df) + sample_proc.invoice_month = "2024-06" + output_df = sample_proc._filter_partners(sample_df) answer_df = pandas.DataFrame({"Institution": ["BU", "HU"]}) self.assertTrue(output_df.equals(answer_df)) # When all partnerships are active - sample_inv.invoice_month = "2024-12" - output_df = sample_inv._filter_partners(sample_df) + sample_proc.invoice_month = "2024-12" + output_df = sample_proc._filter_partners(sample_df) answer_df = pandas.DataFrame({"Institution": ["BU", "HU", "NEU"]}) self.assertTrue(output_df.equals(answer_df)) diff --git a/process_report/tests/util.py b/process_report/tests/util.py index 5145f78..4faf672 100644 --- a/process_report/tests/util.py +++ b/process_report/tests/util.py @@ -12,6 +12,7 @@ validate_pi_alias_processor, lenovo_processor, validate_billable_pi_processor, + new_pi_credit_processor, ) @@ -32,7 +33,7 @@ def new_billable_invoice( nonbillable_pis=None, nonbillable_projects=None, old_pi_filepath="", - limit_new_pi_credit_to_partners=False, + updated_old_pi_df=pandas.DataFrame(), ): if data is None: data = pandas.DataFrame() @@ -45,7 +46,7 @@ def new_billable_invoice( invoice_month, data, old_pi_filepath, - limit_new_pi_credit_to_partners, + updated_old_pi_df, ) @@ -122,3 +123,17 @@ def new_validate_billable_pi_processor( nonbillable_pis, nonbillable_projects, ) + + +def new_new_pi_credit_processor( + name="", + invoice_month="0000-00", + data=None, + old_pi_filepath="", + limit_new_pi_credit_to_partners=False, +): + if data is None: + data = pandas.DataFrame() + return new_pi_credit_processor.NewPICreditProcessor( + name, invoice_month, data, old_pi_filepath, limit_new_pi_credit_to_partners + ) From b9d8086129d0820d8fd003deb4843bdaf85a7210 Mon Sep 17 00:00:00 2001 From: Quan Pham Date: Fri, 8 Nov 2024 14:27:56 -0500 Subject: [PATCH 3/5] Initialized processor for BU subsidy and copied over processing logic --- .../processors/bu_subsidy_processor.py | 62 +++++++++++++++++++ 1 file changed, 62 insertions(+) create mode 100644 process_report/processors/bu_subsidy_processor.py diff --git a/process_report/processors/bu_subsidy_processor.py b/process_report/processors/bu_subsidy_processor.py new file mode 100644 index 0000000..976b8de --- /dev/null +++ b/process_report/processors/bu_subsidy_processor.py @@ -0,0 +1,62 @@ +from dataclasses import dataclass +from decimal import Decimal + +from process_report.invoices import invoice +from process_report.processors import discount_processor + + +@dataclass +class BUSubsidyProcessor(discount_processor.DiscountProcessor): + subsidy_amount: int + + def _prepare(self): + def get_project(row): + project_alloc = row[invoice.PROJECT_FIELD] + if project_alloc.rfind("-") == -1: + return project_alloc + else: + return project_alloc[: project_alloc.rfind("-")] + + self.data = self.data[ + self.data[invoice.IS_BILLABLE_FIELD] & ~self.data[invoice.MISSING_PI_FIELD] + ] + self.data = self.data[ + self.data[invoice.INSTITUTION_FIELD] == "Boston University" + ].copy() + self.data["Project"] = self.data.apply(get_project, axis=1) + self.data[invoice.SUBSIDY_FIELD] = Decimal(0) + + def _process(self): + data_summed_projects = self._sum_project_allocations(self.data) + self.data = self._apply_subsidy(data_summed_projects, self.subsidy_amount) + + def _sum_project_allocations(self, dataframe): + """A project may have multiple allocations, and therefore multiple rows + in the raw invoices. For BU-Internal invoice, we only want 1 row for + each unique project, summing up its allocations' costs""" + project_list = dataframe["Project"].unique() + data_no_dup = dataframe.drop_duplicates("Project", inplace=False) + sum_fields = [invoice.COST_FIELD, invoice.CREDIT_FIELD, invoice.BALANCE_FIELD] + for project in project_list: + project_mask = dataframe["Project"] == project + no_dup_project_mask = data_no_dup["Project"] == project + + sum_fields_sums = dataframe[project_mask][sum_fields].sum().values + data_no_dup.loc[no_dup_project_mask, sum_fields] = sum_fields_sums + + return data_no_dup + + def _apply_subsidy(self, dataframe, subsidy_amount): + pi_list = dataframe[invoice.PI_FIELD].unique() + + for pi in pi_list: + pi_projects = dataframe[dataframe[invoice.PI_FIELD] == pi] + self.apply_flat_discount( + dataframe, + pi_projects, + subsidy_amount, + invoice.SUBSIDY_FIELD, + invoice.BALANCE_FIELD, + ) + + return dataframe From c81622d4049efad96dcb4cd992c564f09b290bad Mon Sep 17 00:00:00 2001 From: Quan Pham Date: Mon, 11 Nov 2024 13:20:42 -0500 Subject: [PATCH 4/5] Since all current discounts have been implemented as processors, the `DiscountInvoice` class is now removed. No invoice class will now class `_prepare()` or `_process()`. `BUSubsidyProcessor`, which handles processing for the BU subsidy, sets `IS_DISCOUNT_BY_NERC` to `False` because the subsidy is not provided by NERC. Because of this, `BU Balance` indicates the money which BU (not the PI they are subsidizing) owes to the MGHPCC. The test cases for the BU Subsidy has been refactored to be more robust and readable. The `Project` field has been added to `invoice.py` --- .../invoices/bu_internal_invoice.py | 48 +--- process_report/invoices/discount_invoice.py | 79 ------ process_report/invoices/invoice.py | 1 + process_report/process_report.py | 9 +- .../processors/bu_subsidy_processor.py | 44 ++- process_report/tests/unit_tests.py | 256 +++++++++++------- process_report/tests/util.py | 25 +- 7 files changed, 211 insertions(+), 251 deletions(-) delete mode 100644 process_report/invoices/discount_invoice.py diff --git a/process_report/invoices/bu_internal_invoice.py b/process_report/invoices/bu_internal_invoice.py index e1f7732..5bd820c 100644 --- a/process_report/invoices/bu_internal_invoice.py +++ b/process_report/invoices/bu_internal_invoice.py @@ -1,12 +1,10 @@ from dataclasses import dataclass -from decimal import Decimal import process_report.invoices.invoice as invoice -import process_report.invoices.discount_invoice as discount_invoice @dataclass -class BUInternalInvoice(discount_invoice.DiscountInvoice): +class BUInternalInvoice(invoice.Invoice): """ This invoice operates on data processed by these Processors: - ValidateBillablePIsProcessor @@ -20,31 +18,19 @@ class BUInternalInvoice(discount_invoice.DiscountInvoice): invoice.COST_FIELD, invoice.CREDIT_FIELD, invoice.SUBSIDY_FIELD, - invoice.BALANCE_FIELD, + invoice.PI_BALANCE_FIELD, ] - subsidy_amount: int - - def _prepare(self): - def get_project(row): - project_alloc = row[invoice.PROJECT_FIELD] - if project_alloc.rfind("-") == -1: - return project_alloc - else: - return project_alloc[: project_alloc.rfind("-")] + exported_columns_map = {invoice.PI_BALANCE_FIELD: "Balance"} + def _prepare_export(self): self.data = self.data[ self.data[invoice.IS_BILLABLE_FIELD] & ~self.data[invoice.MISSING_PI_FIELD] ] self.data = self.data[ self.data[invoice.INSTITUTION_FIELD] == "Boston University" - ].copy() - self.data["Project"] = self.data.apply(get_project, axis=1) - self.data[invoice.SUBSIDY_FIELD] = Decimal(0) - - def _process(self): - data_summed_projects = self._sum_project_allocations(self.data) - self.data = self._apply_subsidy(data_summed_projects, self.subsidy_amount) + ] + self.data = self._sum_project_allocations(self.data) def _sum_project_allocations(self, dataframe): """A project may have multiple allocations, and therefore multiple rows @@ -52,7 +38,12 @@ def _sum_project_allocations(self, dataframe): each unique project, summing up its allocations' costs""" project_list = dataframe["Project"].unique() data_no_dup = dataframe.drop_duplicates("Project", inplace=False) - sum_fields = [invoice.COST_FIELD, invoice.CREDIT_FIELD, invoice.BALANCE_FIELD] + sum_fields = [ + invoice.COST_FIELD, + invoice.CREDIT_FIELD, + invoice.SUBSIDY_FIELD, + invoice.PI_BALANCE_FIELD, + ] for project in project_list: project_mask = dataframe["Project"] == project no_dup_project_mask = data_no_dup["Project"] == project @@ -61,18 +52,3 @@ def _sum_project_allocations(self, dataframe): data_no_dup.loc[no_dup_project_mask, sum_fields] = sum_fields_sums return data_no_dup - - def _apply_subsidy(self, dataframe, subsidy_amount): - pi_list = dataframe[invoice.PI_FIELD].unique() - - for pi in pi_list: - pi_projects = dataframe[dataframe[invoice.PI_FIELD] == pi] - self.apply_flat_discount( - dataframe, - pi_projects, - subsidy_amount, - invoice.SUBSIDY_FIELD, - invoice.BALANCE_FIELD, - ) - - return dataframe diff --git a/process_report/invoices/discount_invoice.py b/process_report/invoices/discount_invoice.py deleted file mode 100644 index 333b998..0000000 --- a/process_report/invoices/discount_invoice.py +++ /dev/null @@ -1,79 +0,0 @@ -from dataclasses import dataclass - -import pandas - -import process_report.invoices.invoice as invoice - - -@dataclass -class DiscountInvoice(invoice.Invoice): - """ - Invoice class containing functions useful for applying discounts - on dataframes - """ - - @staticmethod - def apply_flat_discount( - invoice: pandas.DataFrame, - pi_projects: pandas.DataFrame, - discount_amount: int, - discount_field: str, - balance_field: str, - code_field: str = None, - discount_code: str = None, - ): - """ - Takes in an invoice and a list of PI projects that are a subset of it, - and applies a flat discount to those PI projects. Note that this function - will change the provided `invoice` Dataframe directly. Therefore, it does - not return the changed invoice. - - This function assumes that the balance field shows the remaining cost of the project, - or what the PI would pay before the flat discount is applied. - - If the optional parameters `code_field` and `discount_code` are passed in, - `discount_code` will be comma-APPENDED to the `code_field` of projects where - the discount is applied - - Returns the amount of discount used. - - :param invoice: Dataframe containing all projects - :param pi_projects: A subset of `invoice`, containing all projects for a PI you want to apply the discount - :param discount_amount: The discount given to the PI - :param discount_field: Name of the field to put the discount amount applied to each project - :param balance_field: Name of the balance field - :param code_field: Name of the discount code field - :param discount_code: Code of the discount - """ - - def apply_discount_on_project(remaining_discount_amount, project_i, project): - remaining_project_balance = project[balance_field] - applied_discount = min(remaining_project_balance, remaining_discount_amount) - invoice.at[project_i, discount_field] = applied_discount - invoice.at[project_i, balance_field] = ( - project[balance_field] - applied_discount - ) - remaining_discount_amount -= applied_discount - return remaining_discount_amount - - def apply_credit_code_on_project(project_i): - if code_field and discount_code: - if pandas.isna(invoice.at[project_i, code_field]): - invoice.at[project_i, code_field] = discount_code - else: - invoice.at[project_i, code_field] = ( - invoice.at[project_i, code_field] + "," + discount_code - ) - - remaining_discount_amount = discount_amount - for i, row in pi_projects.iterrows(): - if remaining_discount_amount == 0: - break - else: - remaining_discount_amount = apply_discount_on_project( - remaining_discount_amount, i, row - ) - apply_credit_code_on_project(i) - - discount_used = discount_amount - remaining_discount_amount - return discount_used diff --git a/process_report/invoices/invoice.py b/process_report/invoices/invoice.py index 35c8ffb..7ffb6d8 100644 --- a/process_report/invoices/invoice.py +++ b/process_report/invoices/invoice.py @@ -37,6 +37,7 @@ IS_BILLABLE_FIELD = "Is Billable" MISSING_PI_FIELD = "Missing PI" PI_BALANCE_FIELD = "PI Balance" +PROJECT_NAME_FIELD = "Project" ### diff --git a/process_report/process_report.py b/process_report/process_report.py index cd65430..49b7644 100644 --- a/process_report/process_report.py +++ b/process_report/process_report.py @@ -21,6 +21,7 @@ lenovo_processor, validate_billable_pi_processor, new_pi_credit_processor, + bu_subsidy_processor, ) ### PI file field names @@ -242,7 +243,12 @@ def main(): ) new_pi_credit_proc.process() - processed_data = new_pi_credit_proc.data + bu_subsidy_proc = bu_subsidy_processor.BUSubsidyProcessor( + "", invoice_month, new_pi_credit_proc.data.copy(), args.BU_subsidy_amount + ) + bu_subsidy_proc.process() + + processed_data = bu_subsidy_proc.data ### Initialize invoices @@ -280,7 +286,6 @@ def main(): name=args.BU_invoice_file, invoice_month=invoice_month, data=processed_data.copy(), - subsidy_amount=args.BU_subsidy_amount, ) pi_inv = pi_specific_invoice.PIInvoice( diff --git a/process_report/processors/bu_subsidy_processor.py b/process_report/processors/bu_subsidy_processor.py index 976b8de..8dc35f7 100644 --- a/process_report/processors/bu_subsidy_processor.py +++ b/process_report/processors/bu_subsidy_processor.py @@ -7,6 +7,8 @@ @dataclass class BUSubsidyProcessor(discount_processor.DiscountProcessor): + IS_DISCOUNT_BY_NERC = False + subsidy_amount: int def _prepare(self): @@ -17,43 +19,35 @@ def get_project(row): else: return project_alloc[: project_alloc.rfind("-")] - self.data = self.data[ - self.data[invoice.IS_BILLABLE_FIELD] & ~self.data[invoice.MISSING_PI_FIELD] - ] - self.data = self.data[ - self.data[invoice.INSTITUTION_FIELD] == "Boston University" - ].copy() - self.data["Project"] = self.data.apply(get_project, axis=1) + self.data[invoice.PROJECT_NAME_FIELD] = self.data.apply(get_project, axis=1) self.data[invoice.SUBSIDY_FIELD] = Decimal(0) def _process(self): - data_summed_projects = self._sum_project_allocations(self.data) - self.data = self._apply_subsidy(data_summed_projects, self.subsidy_amount) + self.data = self._apply_subsidy(self.data, self.subsidy_amount) - def _sum_project_allocations(self, dataframe): - """A project may have multiple allocations, and therefore multiple rows - in the raw invoices. For BU-Internal invoice, we only want 1 row for - each unique project, summing up its allocations' costs""" - project_list = dataframe["Project"].unique() - data_no_dup = dataframe.drop_duplicates("Project", inplace=False) - sum_fields = [invoice.COST_FIELD, invoice.CREDIT_FIELD, invoice.BALANCE_FIELD] - for project in project_list: - project_mask = dataframe["Project"] == project - no_dup_project_mask = data_no_dup["Project"] == project - - sum_fields_sums = dataframe[project_mask][sum_fields].sum().values - data_no_dup.loc[no_dup_project_mask, sum_fields] = sum_fields_sums + @staticmethod + def _get_subsidy_eligible_projects(data): + filtered_data = data[ + data[invoice.IS_BILLABLE_FIELD] & ~data[invoice.MISSING_PI_FIELD] + ] + filtered_data = filtered_data[ + filtered_data[invoice.INSTITUTION_FIELD] == "Boston University" + ].copy() - return data_no_dup + return filtered_data def _apply_subsidy(self, dataframe, subsidy_amount): - pi_list = dataframe[invoice.PI_FIELD].unique() + subsidy_eligible_projects = self._get_subsidy_eligible_projects(dataframe) + pi_list = subsidy_eligible_projects[invoice.PI_FIELD].unique() for pi in pi_list: - pi_projects = dataframe[dataframe[invoice.PI_FIELD] == pi] + pi_projects = subsidy_eligible_projects[ + subsidy_eligible_projects[invoice.PI_FIELD] == pi + ] self.apply_flat_discount( dataframe, pi_projects, + invoice.PI_BALANCE_FIELD, subsidy_amount, invoice.SUBSIDY_FIELD, invoice.BALANCE_FIELD, diff --git a/process_report/tests/unit_tests.py b/process_report/tests/unit_tests.py index f09a7de..315ecaf 100644 --- a/process_report/tests/unit_tests.py +++ b/process_report/tests/unit_tests.py @@ -757,119 +757,179 @@ def test_apply_credit_error(self): test_invoice._get_pi_age(old_pi_df, "PI1", invoice_month) -class TestBUSubsidy(TestCase): - def setUp(self): - data = { - "Invoice Month": [ - "2024-03", - "2024-03", - "2024-03", - "2024-03", - "2024-03", - "2024-03", - "2024-03", - "2024-03", - ], - "Manager (PI)": ["PI1", "PI1", "PI2", "PI2", "PI3", "PI3", "PI4", "PI4"], - "Institution": [ - "Boston University", - "Boston University", +class TestBUSubsidyProcessor(TestCase): + def _assert_result_invoice( + self, + subsidy_amount, + test_invoice, + answer_invoice, + invoice_month="0000-00", + ): + new_bu_subsidy_proc = test_utils.new_bu_subsidy_processor( + invoice_month=invoice_month, + data=test_invoice, + subsidy_amount=subsidy_amount, + ) + new_bu_subsidy_proc.process() + output_invoice = new_bu_subsidy_proc.data + answer_invoice = answer_invoice.astype(output_invoice.dtypes) + print(output_invoice) + print(answer_invoice) + + self.assertTrue(output_invoice.equals(answer_invoice)) + + def _get_test_invoice( + self, + pi, + pi_balances, + balances=None, + project_names=None, + institution=None, + is_billable=None, + missing_pi=None, + ): + if not balances: + balances = pi_balances + + if not project_names: + project_names = ["Project" for _ in range(len(pi))] + + if not institution: + institution = ["Boston University" for _ in range(len(pi))] + + if not is_billable: + is_billable = [True for _ in range(len(pi))] + + if not missing_pi: + missing_pi = [False for _ in range(len(pi))] + + return pandas.DataFrame( + { + "Manager (PI)": pi, + "Project - Allocation": project_names, + "PI Balance": pi_balances, + "Balance": balances, + "Institution": institution, + "Is Billable": is_billable, + "Missing PI": missing_pi, + } + ) + + def test_exclude_non_BU_pi(self): + """Are only BU PIs given the subsidy?""" + + subsidy_amount = 100 + test_invoice = self._get_test_invoice( + [str(i) for i in range(5)], + pi_balances=[subsidy_amount for _ in range(5)], + institution=[ "Boston University", "Boston University", - "Harvard University", # Test case for non-BU PIs + "boston university", "Harvard University", - "Boston University", - "Boston University", + "BU", ], - "Project - Allocation": [ - "ProjectA-e6413", - "ProjectA-t575e6", # Test case for project with >1 allocation - "ProjectB-fddgfygg", - "ProjectB-5t143t", - "ProjectC-t14334", - "ProjectD", # Test case for correctly extracting project name - "ProjectE-test-r25135", # Test case for BU PI with >1 project - "ProjectF", - ], - "Cost": [1050, 500, 100, 925, 10000, 1000, 1050, 100], - "Credit": [ - 1000, - 0, - 100, - 900, - 0, - 0, - 1000, - 0, - ], # Test cases where PI does/dones't have credits alreadys - "Balance": [ - 50, - 500, - 0, - 25, - 10000, - 1000, - 50, - 100, - ], # Test case where subsidy does/doesn't cover fully balance - "Is Billable": [True, True, True, True, True, True, True, True], - "Missing PI": [False, False, False, False, False, False, False, False], - } - self.dataframe = pandas.DataFrame(data) - self.subsidy = 100 - - def test_apply_BU_subsidy(self): - test_invoice = test_utils.new_bu_internal_invoice( - data=self.dataframe, subsidy_amount=self.subsidy ) - test_invoice.process() - output_df = test_invoice.data.reset_index() - self.assertTrue( - set( - [ - process_report.INVOICE_DATE_FIELD, - "Project", - process_report.PI_FIELD, - process_report.COST_FIELD, - process_report.CREDIT_FIELD, - process_report.SUBSIDY_FIELD, - process_report.BALANCE_FIELD, - ] - ).issubset(output_df) + answer_invoice = test_invoice.copy() + answer_invoice["Project"] = answer_invoice["Project - Allocation"] + answer_invoice["Subsidy"] = [subsidy_amount, subsidy_amount, 0, 0, 0] + answer_invoice["PI Balance"] = [ + 0, + 0, + subsidy_amount, + subsidy_amount, + subsidy_amount, + ] + + self._assert_result_invoice(subsidy_amount, test_invoice, answer_invoice) + + def test_exclude_nonbillables(self): + """Are nonbillables excluded from the subsidy?""" + subsidy_amount = 100 + test_invoice = self._get_test_invoice( + [str(i) for i in range(6)], + pi_balances=[subsidy_amount for _ in range(6)], + is_billable=[True, True, False, False, True, True], + missing_pi=[True, True, False, False, False, False], ) - self.assertTrue( - set(["PI1", "PI2", "PI4"]).issubset(output_df["Manager (PI)"].unique()) + answer_invoice = test_invoice.copy() + answer_invoice["Project"] = answer_invoice["Project - Allocation"] + answer_invoice["Subsidy"] = [0, 0, 0, 0, subsidy_amount, subsidy_amount] + answer_invoice["PI Balance"] = [ + subsidy_amount, + subsidy_amount, + subsidy_amount, + subsidy_amount, + 0, + 0, + ] + + self._assert_result_invoice(subsidy_amount, test_invoice, answer_invoice) + + def test_one_pi_many_allocations(self): + """Is subsidy applied properly to BU PI with many allocations?""" + + # Two projects, one allocation each + subsidy_amount = 100 + test_invoice = self._get_test_invoice( + ["PI" for i in range(2)], + pi_balances=[60, 60], + project_names=["P1", "P2"], ) - self.assertFalse("PI3" in output_df["Project"].unique()) - self.assertTrue( - set(["ProjectA", "ProjectB", "ProjectE-test", "ProjectF"]).issubset( - output_df["Project"].unique() - ) + answer_invoice = test_invoice.copy() + answer_invoice["Project"] = answer_invoice["Project - Allocation"] + answer_invoice["Subsidy"] = [60, 40] + answer_invoice["PI Balance"] = [0, 20] + + self._assert_result_invoice(subsidy_amount, test_invoice, answer_invoice) + + # Two projects, two allocations each + test_invoice = self._get_test_invoice( + ["PI" for i in range(4)], + pi_balances=[40, 40, 40, 40], + project_names=["P1-A1", "P1-A1-test", "P2", "P2-"], ) - self.assertFalse( - set(["ProjectC-t14334", "ProjectC", "ProjectD"]).intersection( - output_df["Project"].unique() - ) + + answer_invoice = test_invoice.copy() + answer_invoice["Project"] = ["P1", "P1-A1", "P2", "P2"] + answer_invoice["Subsidy"] = [40, 40, 20, 0] + answer_invoice["PI Balance"] = [0, 0, 20, 40] + + self._assert_result_invoice(subsidy_amount, test_invoice, answer_invoice) + + # Two allocations, one where PI balance != NERC balance + test_invoice = self._get_test_invoice( + ["PI" for i in range(2)], + pi_balances=[80, 80], + project_names=["P1", "P2"], + balances=[100, 80], ) - self.assertEqual(4, len(output_df.index)) - self.assertEqual(1550, output_df.loc[0, "Cost"]) - self.assertEqual(1025, output_df.loc[1, "Cost"]) - self.assertEqual(1050, output_df.loc[2, "Cost"]) - self.assertEqual(100, output_df.loc[3, "Cost"]) + answer_invoice = test_invoice.copy() + answer_invoice["Project"] = answer_invoice["Project - Allocation"] + answer_invoice["Subsidy"] = [80, 20] + answer_invoice["PI Balance"] = [0, 60] + + self._assert_result_invoice(subsidy_amount, test_invoice, answer_invoice) - self.assertEqual(100, output_df.loc[0, "Subsidy"]) - self.assertEqual(25, output_df.loc[1, "Subsidy"]) - self.assertEqual(50, output_df.loc[2, "Subsidy"]) - self.assertEqual(50, output_df.loc[3, "Subsidy"]) + def test_two_pi(self): + """Is subsidy applied to more than one PI?""" + # Each PI has two allocations + subsidy_amount = 100 + test_invoice = self._get_test_invoice( + ["PI1", "PI1", "PI2", "PI2"], + pi_balances=[80, 80, 40, 40], + ) + + answer_invoice = test_invoice.copy() + answer_invoice["Project"] = answer_invoice["Project - Allocation"] + answer_invoice["Subsidy"] = [80, 20, 40, 40] + answer_invoice["PI Balance"] = [0, 60, 0, 0] - self.assertEqual(450, output_df.loc[0, "Balance"]) - self.assertEqual(0, output_df.loc[1, "Balance"]) - self.assertEqual(0, output_df.loc[2, "Balance"]) - self.assertEqual(50, output_df.loc[3, "Balance"]) + self._assert_result_invoice(subsidy_amount, test_invoice, answer_invoice) class TestLenovoProcessor(TestCase): diff --git a/process_report/tests/util.py b/process_report/tests/util.py index 4faf672..cfe0bbb 100644 --- a/process_report/tests/util.py +++ b/process_report/tests/util.py @@ -3,7 +3,6 @@ from process_report.invoices import ( invoice, billable_invoice, - bu_internal_invoice, pi_specific_invoice, ) @@ -13,6 +12,7 @@ lenovo_processor, validate_billable_pi_processor, new_pi_credit_processor, + bu_subsidy_processor, ) @@ -50,16 +50,6 @@ def new_billable_invoice( ) -def new_bu_internal_invoice( - name="", invoice_month="0000-00", data=None, subsidy_amount=0 -): - if data is None: - data = pandas.DataFrame() - return bu_internal_invoice.BUInternalInvoice( - name, invoice_month, data, subsidy_amount - ) - - def new_pi_specific_invoice( name="", invoice_month="0000-00", @@ -137,3 +127,16 @@ def new_new_pi_credit_processor( return new_pi_credit_processor.NewPICreditProcessor( name, invoice_month, data, old_pi_filepath, limit_new_pi_credit_to_partners ) + + +def new_bu_subsidy_processor( + name="", + invoice_month="0000-00", + data=None, + subsidy_amount=0, +): + if data is None: + data = pandas.DataFrame() + return bu_subsidy_processor.BUSubsidyProcessor( + name, invoice_month, data, subsidy_amount + ) From 23b3a8ae0100aa5266cbecf1e509a7270a7eb91b Mon Sep 17 00:00:00 2001 From: Quan Pham Date: Wed, 4 Dec 2024 11:32:15 -0500 Subject: [PATCH 5/5] Invoices no longer assign to `self.data` during filtering step Instead, all invoice classes will now assign to `self.export_data` This removes ambiguity on whether the invoices will modify their internal dataframe and the need to call `copy()` on the processed dataframe --- process_report/invoices/NERC_total_invoice.py | 6 +++--- process_report/invoices/billable_invoice.py | 2 +- process_report/invoices/bu_internal_invoice.py | 8 ++++---- process_report/invoices/invoice.py | 7 ++++--- process_report/invoices/lenovo_invoice.py | 2 +- process_report/invoices/nonbillable_invoice.py | 2 +- process_report/invoices/pi_specific_invoice.py | 4 ++-- process_report/process_report.py | 12 ++++++------ process_report/tests/unit_tests.py | 6 ++++-- 9 files changed, 26 insertions(+), 23 deletions(-) diff --git a/process_report/invoices/NERC_total_invoice.py b/process_report/invoices/NERC_total_invoice.py index 92982da..e39e49b 100644 --- a/process_report/invoices/NERC_total_invoice.py +++ b/process_report/invoices/NERC_total_invoice.py @@ -51,9 +51,9 @@ def output_s3_archive_key(self): return f"Invoices/{self.invoice_month}/Archive/NERC-{self.invoice_month}-Total-Invoice {util.get_iso8601_time()}.csv" def _prepare_export(self): - self.data = self.data[ + self.export_data = self.data[ self.data[invoice.IS_BILLABLE_FIELD] & ~self.data[invoice.MISSING_PI_FIELD] ] - self.data = self.data[ - self.data[invoice.INSTITUTION_FIELD].isin(self.INCLUDED_INSTITUTIONS) + self.export_data = self.export_data[ + self.export_data[invoice.INSTITUTION_FIELD].isin(self.INCLUDED_INSTITUTIONS) ].copy() diff --git a/process_report/invoices/billable_invoice.py b/process_report/invoices/billable_invoice.py index 1761cc8..8bab38a 100644 --- a/process_report/invoices/billable_invoice.py +++ b/process_report/invoices/billable_invoice.py @@ -42,7 +42,7 @@ class BillableInvoice(invoice.Invoice): ] def _prepare_export(self): - self.data = self.data[ + self.export_data = self.data[ self.data[invoice.IS_BILLABLE_FIELD] & ~self.data[invoice.MISSING_PI_FIELD] ] self.updated_old_pi_df = self.updated_old_pi_df.astype( diff --git a/process_report/invoices/bu_internal_invoice.py b/process_report/invoices/bu_internal_invoice.py index 5bd820c..0db03b9 100644 --- a/process_report/invoices/bu_internal_invoice.py +++ b/process_report/invoices/bu_internal_invoice.py @@ -24,13 +24,13 @@ class BUInternalInvoice(invoice.Invoice): exported_columns_map = {invoice.PI_BALANCE_FIELD: "Balance"} def _prepare_export(self): - self.data = self.data[ + self.export_data = self.data[ self.data[invoice.IS_BILLABLE_FIELD] & ~self.data[invoice.MISSING_PI_FIELD] ] - self.data = self.data[ - self.data[invoice.INSTITUTION_FIELD] == "Boston University" + self.export_data = self.export_data[ + self.export_data[invoice.INSTITUTION_FIELD] == "Boston University" ] - self.data = self._sum_project_allocations(self.data) + self.export_data = self._sum_project_allocations(self.export_data) def _sum_project_allocations(self, dataframe): """A project may have multiple allocations, and therefore multiple rows diff --git a/process_report/invoices/invoice.py b/process_report/invoices/invoice.py index 7ffb6d8..47797c8 100644 --- a/process_report/invoices/invoice.py +++ b/process_report/invoices/invoice.py @@ -49,6 +49,7 @@ class Invoice: name: str invoice_month: str data: pandas.DataFrame + export_data = None def process(self): self._prepare() @@ -93,13 +94,13 @@ def _prepare_export(self): def _filter_columns(self): """Filters and renames columns before exporting""" - return self.data.copy()[self.export_columns_list].rename( + self.export_data = self.export_data[self.export_columns_list].rename( columns=self.exported_columns_map ) def export(self): - export_data = self._filter_columns() - export_data.to_csv(self.output_path, index=False) + self._filter_columns() + self.export_data.to_csv(self.output_path, index=False) def export_s3(self, s3_bucket): s3_bucket.upload_file(self.output_path, self.output_s3_key) diff --git a/process_report/invoices/lenovo_invoice.py b/process_report/invoices/lenovo_invoice.py index 1fbb5c7..fe96973 100644 --- a/process_report/invoices/lenovo_invoice.py +++ b/process_report/invoices/lenovo_invoice.py @@ -19,6 +19,6 @@ class LenovoInvoice(invoice.Invoice): exported_columns_map = {invoice.SU_HOURS_FIELD: "SU Hours"} def _prepare_export(self): - self.data = self.data[ + self.export_data = self.data[ self.data[invoice.SU_TYPE_FIELD].isin(self.LENOVO_SU_TYPES) ] diff --git a/process_report/invoices/nonbillable_invoice.py b/process_report/invoices/nonbillable_invoice.py index 578e1d2..350e290 100644 --- a/process_report/invoices/nonbillable_invoice.py +++ b/process_report/invoices/nonbillable_invoice.py @@ -24,4 +24,4 @@ class NonbillableInvoice(invoice.Invoice): ] def _prepare_export(self): - self.data = self.data[~self.data[invoice.IS_BILLABLE_FIELD]] + self.export_data = self.data[~self.data[invoice.IS_BILLABLE_FIELD]] diff --git a/process_report/invoices/pi_specific_invoice.py b/process_report/invoices/pi_specific_invoice.py index dabc0a0..5fc2e1a 100644 --- a/process_report/invoices/pi_specific_invoice.py +++ b/process_report/invoices/pi_specific_invoice.py @@ -34,10 +34,10 @@ class PIInvoice(invoice.Invoice): ] def _prepare(self): - self.data = self.data[ + self.export_data = self.data[ self.data[invoice.IS_BILLABLE_FIELD] & ~self.data[invoice.MISSING_PI_FIELD] ] - self.pi_list = self.data[invoice.PI_FIELD].unique() + self.pi_list = self.export_data[invoice.PI_FIELD].unique() def export(self): def _export_pi_invoice(pi): diff --git a/process_report/process_report.py b/process_report/process_report.py index 49b7644..5a1ed77 100644 --- a/process_report/process_report.py +++ b/process_report/process_report.py @@ -255,12 +255,12 @@ def main(): lenovo_inv = lenovo_invoice.LenovoInvoice( name=args.Lenovo_file, invoice_month=invoice_month, - data=processed_data.copy(), + data=processed_data, ) nonbillable_inv = nonbillable_invoice.NonbillableInvoice( name=args.nonbillable_file, invoice_month=invoice_month, - data=processed_data.copy(), + data=processed_data, nonbillable_pis=pi, nonbillable_projects=projects, ) @@ -271,7 +271,7 @@ def main(): billable_inv = billable_invoice.BillableInvoice( name=args.output_file, invoice_month=invoice_month, - data=processed_data.copy(), + data=processed_data, old_pi_filepath=old_pi_file, updated_old_pi_df=new_pi_credit_proc.updated_old_pi_df, ) @@ -279,17 +279,17 @@ def main(): nerc_total_inv = NERC_total_invoice.NERCTotalInvoice( name=args.NERC_total_invoice_file, invoice_month=invoice_month, - data=processed_data.copy(), + data=processed_data, ) bu_internal_inv = bu_internal_invoice.BUInternalInvoice( name=args.BU_invoice_file, invoice_month=invoice_month, - data=processed_data.copy(), + data=processed_data, ) pi_inv = pi_specific_invoice.PIInvoice( - name=args.output_folder, invoice_month=invoice_month, data=processed_data.copy() + name=args.output_folder, invoice_month=invoice_month, data=processed_data ) util.process_and_export_invoices( diff --git a/process_report/tests/unit_tests.py b/process_report/tests/unit_tests.py index 315ecaf..6ef99e6 100644 --- a/process_report/tests/unit_tests.py +++ b/process_report/tests/unit_tests.py @@ -1036,9 +1036,11 @@ class TestBaseInvoice(TestCase): def test_filter_exported_columns(self): test_invoice = pandas.DataFrame(columns=["C1", "C2", "C3", "C4", "C5"]) answer_invoice = pandas.DataFrame(columns=["C1", "C3R", "C5R"]) - inv = test_utils.new_base_invoice(data=test_invoice) + inv = test_utils.new_base_invoice() + inv.export_data = test_invoice inv.export_columns_list = ["C1", "C3", "C5"] inv.exported_columns_map = {"C3": "C3R", "C5": "C5R"} - result_invoice = inv._filter_columns() + inv._filter_columns() + result_invoice = inv.export_data self.assertTrue(result_invoice.equals(answer_invoice))