From 98b4b68d22bb1dad4b57748f11800f059e6bd64a Mon Sep 17 00:00:00 2001
From: QuanMPhm <qmpham2019@gmail.com>
Date: Thu, 4 Apr 2024 14:00:11 -0400
Subject: [PATCH] Added processing and exporting for Lenovo SU Types

---
 process_report/process_report.py   | 21 +++++++++++++++++
 process_report/tests/unit_tests.py | 36 ++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/process_report/process_report.py b/process_report/process_report.py
index 073ed72..5dbb246 100644
--- a/process_report/process_report.py
+++ b/process_report/process_report.py
@@ -171,6 +171,7 @@ def main():
     export_pi_billables(billable_projects, args.output_folder)
     export_HU_only(billable_projects, args.HU_only)
     export_HU_BU(billable_projects, args.HU_BU)
+    export_lenovo(billable_projects)
 
 
 def merge_csv(files):
@@ -275,5 +276,25 @@ def export_HU_BU(dataframe, output_file):
     HU_BU_projects.to_csv(output_file)        
 
 
+def export_lenovo(dataframe: pandas.DataFrame, output_file=None):
+
+    lenovo_file_name = output_file or f'Lenovo_{dataframe[INVOICE_DATE_FIELD].iat[0]}.csv'
+
+    LENOVO_SU_TYPES = ['OpenShift GPUA100SXM4', 'OpenStack GPUA100SXM4']
+    SU_CHARGE_MULTIPLIER = 1
+
+    lenovo_df = dataframe[dataframe[SU_TYPE_FIELD].isin(LENOVO_SU_TYPES)][[
+            INVOICE_DATE_FIELD, 
+            PROJECT_FIELD, 
+            INSTITUTION_FIELD, 
+            SU_HOURS_FIELD, 
+            SU_TYPE_FIELD]]
+    
+    lenovo_df.rename(columns={SU_HOURS_FIELD: 'SU Hours'}, inplace=True)
+    lenovo_df.insert(len(lenovo_df.columns), 'SU Charge', SU_CHARGE_MULTIPLIER)
+    lenovo_df['Charge'] = lenovo_df['SU Hours'] * lenovo_df['SU Charge']
+    lenovo_df.to_csv(lenovo_file_name)
+
+
 if __name__ == "__main__":
     main()
diff --git a/process_report/tests/unit_tests.py b/process_report/tests/unit_tests.py
index 90b1473..3953ae9 100644
--- a/process_report/tests/unit_tests.py
+++ b/process_report/tests/unit_tests.py
@@ -237,4 +237,40 @@ def test_apply_credit_0002(self):
         self.assertEqual(0, credited_projects[credited_projects['Project - Allocation'] == 'ProjectE']['Balance'].iloc[0])
         self.assertEqual(800, credited_projects[credited_projects['Project - Allocation'] == 'ProjectF']['Balance'].iloc[0])
 
+
+class TestExportLenovo(TestCase):
+    def setUp(self):
+
+        data = {
+            'Invoice Month': ['2023-01','2023-01','2023-01','2023-01','2023-01', '2023-01'],
+            'Project - Allocation': ['ProjectA', 'ProjectB', 'ProjectC', 'ProjectD', 'ProjectE', 'ProjectF'],
+            'Institution': ['A', 'B', 'C', 'D', 'E', 'F'],
+            'SU Hours (GBhr or SUhr)': [1, 10, 100, 4, 432, 10],
+            'SU Type': ['OpenShift GPUA100SXM4', 'CPU', 'OpenShift GPUA100SXM4', 'OpenStack GPUA100SXM4', 'CPU', 'GPU']
+        }
+        self.dataframe = pandas.DataFrame(data)
+
+        output_file = tempfile.NamedTemporaryFile(delete=False, mode='w', suffix='.csv')
+        self.output_file = output_file.name
+
+    def tearDown(self):
+        os.remove(self.output_file)
+
+
+    def test_apply_credit_0002(self):
+        process_report.export_lenovo(self.dataframe, self.output_file)
+        output_df = pandas.read_csv(self.output_file)
+
+        self.assertTrue(set([
+            process_report.INVOICE_MONTH_FIELD, 
+            process_report.PROJECT_FIELD, 
+            process_report.INSTITUTION_FIELD, 
+            process_report.SU_TYPE_FIELD,
+            'SU Hours', 
+            'SU Charge',
+            'Charge',
+        ]).issubset(output_df))
         
+        for i, row in output_df.iterrows():
+            self.assertIn(row[process_report.SU_TYPE_FIELD], ['OpenShift GPUA100SXM4', 'OpenStack GPUA100SXM4'])
+            self.assertEqual(row['Charge'], row['SU Charge'] * row['SU Hours'])