Refactored PI-specific invoices

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

@@ -1,5 +1,4 @@
 import argparse
-import os
 import sys
 import datetime
 
@@ -14,6 +13,7 @@
     billable_invoice,
     NERC_total_invoice,
     bu_internal_invoice,
+    pi_specific_invoice,
 )
 
 
@@ -262,17 +262,13 @@ def main():
         subsidy_amount=args.BU_subsidy_amount,
     )
 
-    process_and_export_invoices([nerc_total_inv, bu_internal_inv], args.upload_to_s3)
-
-    export_pi_billables(billable_inv.data.copy(), args.output_folder, invoice_month)
-
-    if args.upload_to_s3:
-        invoice_list = list()
-
-        for pi_invoice in os.listdir(args.output_folder):
-            invoice_list.append(os.path.join(args.output_folder, pi_invoice))
+    pi_inv = pi_specific_invoice.PIInvoice(
+        name=args.output_folder, invoice_month=invoice_month, data=billable_inv.data
+    )
 
-        upload_to_s3(invoice_list, invoice_month)
+    process_and_export_invoices(
+        [nerc_total_inv, bu_internal_inv, pi_inv], args.upload_to_s3
+    )
 
 
 def fetch_s3_invoices(invoice_month):
@@ -390,33 +386,5 @@ 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 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

@@ -182,12 +182,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(
@@ -789,32 +790,46 @@ def test_process_lenovo(self):
 
 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.patch("process_report.util.get_iso8601_time")
+    def test_upload_to_s3(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"]
+        filenames = ["test-test", "test2.test", "test3"]
+        sample_base_invoice = test_utils.new_base_invoice(invoice_month=invoice_month)
+
         answers = [
-            ("test.csv", f"Invoices/{invoice_month}/test {invoice_month}.csv"),
             (
-                "test.csv",
-                f"Invoices/{invoice_month}/Archive/test {invoice_month} 0.csv",
+                f"test-test {invoice_month}.csv",
+                f"Invoices/{invoice_month}/test-test {invoice_month}.csv",
             ),
             (
-                "test2.test.csv",
+                f"test-test {invoice_month}.csv",
+                f"Invoices/{invoice_month}/Archive/test-test {invoice_month} 0.csv",
+            ),
+            (
+                f"test2.test {invoice_month}.csv",
                 f"Invoices/{invoice_month}/test2.test {invoice_month}.csv",
             ),
             (
-                "test2.test.csv",
+                f"test2.test {invoice_month}.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"),
+            (
+                f"test3 {invoice_month}.csv",
+                f"Invoices/{invoice_month}/test3 {invoice_month}.csv",
+            ),
+            (
+                f"test3 {invoice_month}.csv",
+                f"Invoices/{invoice_month}/Archive/test3 {invoice_month} 0.csv",
+            ),
         ]
 
-        process_report.upload_to_s3(filenames, invoice_month)
+        for filename in filenames:
+            sample_base_invoice.name = filename
+            sample_base_invoice.export_s3(mock_bucket)
+
         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,19 @@
 import pandas
 
-from process_report.invoices import billable_invoice, bu_internal_invoice
+from process_report.invoices import (
+    invoice,
+    billable_invoice,
+    bu_internal_invoice,
+    pi_specific_invoice,
+)
+
+
+def new_base_invoice(
+    name="",
+    invoice_month="0000-00",
+    data=pandas.DataFrame(),
+):
+    return invoice.Invoice(name, invoice_month, data)
 
 
 def new_billable_invoice(
@@ -27,3 +40,15 @@ def new_bu_internal_invoice(
     return bu_internal_invoice.BUInternalInvoice(
         name, invoice_month, data, subsidy_amount
     )
+
+
+def new_pi_specific_invoice(
+    name="",
+    invoice_month="0000-00",
+    data=pandas.DataFrame(),
+):
+    return pi_specific_invoice.PIInvoice(
+        name,
+        invoice_month,
+        data,
+    )