Refactored PI-specific invoices

The PI-specific is now implemented in the `PIInvoice` class.
This class takes in the entire billable invoice and will
export the PI invoices to local storage and S3

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)

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()

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)

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,
+    )