Merge pull request #14 from QuanMPhm/13/lenovo_only

Added processing and exporting for Lenovo SU Types
CCI-MOC · Apr 12, 2024 · 885761c · 885761c
2 parents 7ba0ee0 + aec6733
commit 885761c
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 0 deletions.
diff --git a/process_report/process_report.py b/process_report/process_report.py
@@ -145,6 +145,7 @@ def main():
     export_pi_billables(billable_projects, args.output_folder)
     export_HU_only(billable_projects, args.HU_invoice_file)
     export_HU_BU(billable_projects, args.HU_BU_invoice_file)
+    export_lenovo(billable_projects)
 
 
 def merge_csv(files):
@@ -295,5 +296,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
@@ -274,3 +274,41 @@ def test_validate_billables(self):
         self.assertEqual(1, len(self.dataframe[pandas.isna(self.dataframe['Manager (PI)'])]))
         validated_df = process_report.validate_pi_names(self.dataframe)
         self.assertEqual(0, len(validated_df[pandas.isna(validated_df['Manager (PI)'])]))
+
+
+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', 'OpenShift GPUA100', 'OpenShift GPUA100SXM4', 'OpenStack GPUA100SXM4', 'OpenStack CPU', 'OpenStack GPUK80']
+        }
+        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_DATE_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'])