diff --git a/process_report/invoices/pi_specific_invoice.py b/process_report/invoices/pi_specific_invoice.py new file mode 100644 index 0000000..fc2bd63 --- /dev/null +++ b/process_report/invoices/pi_specific_invoice.py @@ -0,0 +1,43 @@ +import os +from dataclasses import dataclass + +import pandas + +import process_report.invoices.invoice as invoice +import process_report.util as util + + +@dataclass +class PIInvoice(invoice.Invoice): + def _prepare(self): + self.pi_list = self.data[invoice.PI_FIELD].unique() + + def export(self): + def _export_pi_invoice(pi): + if pandas.isna(pi): + return + pi_projects = self.data[self.data[invoice.PI_FIELD] == pi] + pi_instituition = pi_projects[invoice.INSTITUTION_FIELD].iat[0] + pi_projects.to_csv( + f"{self.name}/{pi_instituition}_{pi} {self.invoice_month}.csv" + ) + + if not os.path.exists( + self.name + ): # self.name is name of folder storing invoices + os.mkdir(self.name) + + for pi in self.pi_list: + _export_pi_invoice(pi) + + def export_s3(self, s3_bucket): + def _export_s3_pi_invoice(pi_invoice): + pi_invoice_path = os.path.join(self.name, pi_invoice) + striped_invoice_path = os.path.splitext(pi_invoice_path)[0] + output_s3_path = f"Invoices/{self.invoice_month}/{striped_invoice_path}.csv" + output_s3_archive_path = f"Invoices/{self.invoice_month}/Archive/{striped_invoice_path} {util.get_iso8601_time()}.csv" + s3_bucket.upload_file(pi_invoice_path, output_s3_path) + s3_bucket.upload_file(pi_invoice_path, output_s3_archive_path) + + for pi_invoice in os.listdir(self.name): + _export_s3_pi_invoice(pi_invoice) diff --git a/process_report/process_report.py b/process_report/process_report.py index 8af2545..f60e029 100644 --- a/process_report/process_report.py +++ b/process_report/process_report.py @@ -15,6 +15,7 @@ nonbillable_invoice, billable_invoice, hu_bu_invoice, + pi_specific_invoice, ) @@ -283,16 +284,16 @@ def main(): bucket = get_invoice_bucket() invoice.export_s3(bucket) - export_pi_billables(billable_inv.data, args.output_folder, invoice_month) - export_BU_only(billable_inv.data, args.BU_invoice_file, args.BU_subsidy_amount) - + pi_inv = pi_specific_invoice.PIInvoice( + name=args.output_folder, invoice_month=invoice_month, data=billable_inv.data + ) + pi_inv.process() + pi_inv.export() if args.upload_to_s3: - invoice_list = list() + bucket = get_invoice_bucket() + pi_inv.export_s3(bucket) - for pi_invoice in os.listdir(args.output_folder): - invoice_list.append(os.path.join(args.output_folder, pi_invoice)) - - upload_to_s3(invoice_list, invoice_month) + export_BU_only(billable_inv.data, args.BU_invoice_file, args.BU_subsidy_amount) def fetch_s3_invoices(invoice_month): @@ -410,22 +411,6 @@ def export_billables(dataframe, output_file): dataframe.to_csv(output_file, index=False) -def export_pi_billables(dataframe: pandas.DataFrame, output_folder, invoice_month): - if not os.path.exists(output_folder): - os.mkdir(output_folder) - - pi_list = dataframe[PI_FIELD].unique() - - for pi in pi_list: - if pandas.isna(pi): - continue - pi_projects = dataframe[dataframe[PI_FIELD] == pi] - pi_instituition = pi_projects[INSTITUTION_FIELD].iat[0] - pi_projects.to_csv( - output_folder + f"/{pi_instituition}_{pi}_{invoice_month}.csv", index=False - ) - - def export_BU_only(dataframe: pandas.DataFrame, output_file, subsidy_amount): def get_project(row): project_alloc = row[PROJECT_FIELD] @@ -483,17 +468,5 @@ def _apply_subsidy(dataframe, subsidy_amount): return dataframe -def upload_to_s3(invoice_list: list, invoice_month): - invoice_bucket = get_invoice_bucket() - for invoice_filename in invoice_list: - striped_filename = os.path.splitext(invoice_filename)[0] - invoice_s3_path = ( - f"Invoices/{invoice_month}/{striped_filename} {invoice_month}.csv" - ) - invoice_s3_path_archive = f"Invoices/{invoice_month}/Archive/{striped_filename} {invoice_month} {get_iso8601_time()}.csv" - invoice_bucket.upload_file(invoice_filename, invoice_s3_path) - invoice_bucket.upload_file(invoice_filename, invoice_s3_path_archive) - - if __name__ == "__main__": main() diff --git a/process_report/tests/unit_tests.py b/process_report/tests/unit_tests.py index 117412d..0b67593 100644 --- a/process_report/tests/unit_tests.py +++ b/process_report/tests/unit_tests.py @@ -1,4 +1,4 @@ -from unittest import TestCase, mock +from unittest import TestCase import tempfile import pandas import pyarrow @@ -183,12 +183,13 @@ def setUp(self): def test_export_pi(self): output_dir = tempfile.TemporaryDirectory() - process_report.export_pi_billables( - self.dataframe, output_dir.name, self.invoice_month + pi_inv = test_utils.new_pi_specific_invoice( + output_dir.name, invoice_month=self.invoice_month, data=self.dataframe ) - - pi_csv_1 = f'{self.dataframe["Institution"][0]}_{self.dataframe["Manager (PI)"][0]}_{self.dataframe["Invoice Month"][0]}.csv' - pi_csv_2 = f'{self.dataframe["Institution"][3]}_{self.dataframe["Manager (PI)"][3]}_{self.dataframe["Invoice Month"][3]}.csv' + pi_inv.process() + pi_inv.export() + pi_csv_1 = f'{self.dataframe["Institution"][0]}_{self.dataframe["Manager (PI)"][0]} {self.dataframe["Invoice Month"][0]}.csv' + pi_csv_2 = f'{self.dataframe["Institution"][3]}_{self.dataframe["Manager (PI)"][3]} {self.dataframe["Invoice Month"][3]}.csv' self.assertIn(pi_csv_1, os.listdir(output_dir.name)) self.assertIn(pi_csv_2, os.listdir(output_dir.name)) self.assertEqual( @@ -785,36 +786,3 @@ def test_process_lenovo(self): ["OpenShift GPUA100SXM4", "OpenStack GPUA100SXM4"], ) self.assertEqual(row["Charge"], row["SU Charge"] * row["SU Hours"]) - - -class TestUploadToS3(TestCase): - @mock.patch("process_report.process_report.get_invoice_bucket") - @mock.patch("process_report.process_report.get_iso8601_time") - def test_remove_prefix(self, mock_get_time, mock_get_bucket): - mock_bucket = mock.MagicMock() - mock_get_bucket.return_value = mock_bucket - mock_get_time.return_value = "0" - - invoice_month = "2024-03" - filenames = ["test.csv", "test2.test.csv", "test3"] - answers = [ - ("test.csv", f"Invoices/{invoice_month}/test {invoice_month}.csv"), - ( - "test.csv", - f"Invoices/{invoice_month}/Archive/test {invoice_month} 0.csv", - ), - ( - "test2.test.csv", - f"Invoices/{invoice_month}/test2.test {invoice_month}.csv", - ), - ( - "test2.test.csv", - f"Invoices/{invoice_month}/Archive/test2.test {invoice_month} 0.csv", - ), - ("test3", f"Invoices/{invoice_month}/test3 {invoice_month}.csv"), - ("test3", f"Invoices/{invoice_month}/Archive/test3 {invoice_month} 0.csv"), - ] - - process_report.upload_to_s3(filenames, invoice_month) - for i, call_args in enumerate(mock_bucket.upload_file.call_args_list): - self.assertTrue(answers[i] in call_args) diff --git a/process_report/tests/util.py b/process_report/tests/util.py index 2a04251..c496064 100644 --- a/process_report/tests/util.py +++ b/process_report/tests/util.py @@ -1,6 +1,6 @@ import pandas -from process_report.invoices import billable_invoice +from process_report.invoices import billable_invoice, pi_specific_invoice def new_billable_invoice( @@ -19,3 +19,15 @@ def new_billable_invoice( nonbillable_projects, old_pi_filepath, ) + + +def new_pi_specific_invoice( + name="", + invoice_month="0000-00", + data=pandas.DataFrame(), +): + return pi_specific_invoice.PIInvoice( + name, + invoice_month, + data, + )