From dc662061d3aeb0b1979a3e645bc43a339bf560a1 Mon Sep 17 00:00:00 2001
From: Quan Pham <qmpham2019@gmail.com>
Date: Thu, 19 Dec 2024 13:33:33 -0500
Subject: [PATCH] NERC Prepay Credits file is now exported through
 `PrepaymentProcessor`

An optional argument, `export_NERC_credits` is added for testing
purposes, to prevent the credits files from being written
while running test cases
---
 .../processors/prepayment_processor.py        | 46 +++++++++++++++++++
 .../processors/test_prepayment_processor.py   | 23 ++++++++++
 process_report/tests/util.py                  |  1 +
 3 files changed, 70 insertions(+)

diff --git a/process_report/processors/prepayment_processor.py b/process_report/processors/prepayment_processor.py
index e8ee863..089fca8 100644
--- a/process_report/processors/prepayment_processor.py
+++ b/process_report/processors/prepayment_processor.py
@@ -22,11 +22,24 @@ class PrepaymentProcessor(discount_processor.DiscountProcessor):
     def PREPAY_DEBITS_S3_BACKUP_FILEPATH(self):
         return f"Prepay/Archive/prepay_debits {util.get_iso8601_time()}.csv"
 
+    @property
+    def CREDITS_SNAPSHOT_FILEPATH(self):
+        return f"NERC_Prepaid_Group-Credits-{self.invoice_month}.csv"
+
+    @property
+    def CREDITS_SNAPSHOT_S3_FILEPATH(self):
+        return f"Invoices/{self.invoice_month}/NERC_Prepaid_Group-Credits-{self.invoice_month}.csv"
+
+    @property
+    def CREDITS_SNAPSHOT_S3_ARCHIVE_FILEPATH(self):
+        return f"Invoices/{self.invoice_month}/Archive/NERC_Prepaid_Group-Credits-{self.invoice_month} {util.get_iso8601_time()}.csv"
+
     prepay_credits: pandas.DataFrame
     prepay_projects: pandas.DataFrame
     prepay_contacts: pandas.DataFrame
     prepay_debits_filepath: str
     upload_to_s3: bool
+    export_NERC_credits: bool = True  # For testing purposes
 
     @staticmethod
     def _load_prepay_debits(prepay_debits_filepath):
@@ -53,6 +66,12 @@ def _process(self):
         self._add_prepay_info()
         self._apply_prepayments()
 
+        if self.export_NERC_credits:
+            credits_snapshot = self._get_prepay_credits_snapshot()
+            self._export_prepay_credits_snapshot(credits_snapshot)
+            if self.upload_to_s3:
+                self._export_s3_prepay_credits_snapshot()
+
         self._export_prepay_debits()
         if self.upload_to_s3:
             self._export_s3_prepay_debits()
@@ -201,12 +220,39 @@ def _apply_prepayments(self):
                         debit_entry_mask, invoice.PREPAY_DEBIT_FIELD
                     ] = prepay_amount_used
 
+    def _get_prepay_credits_snapshot(self):
+        managed_groups_list = list()
+        for group_name, group_dict in self.group_info_dict.items():
+            if group_dict[invoice.PREPAY_MANAGED_FIELD]:
+                managed_groups_list.append(group_name)
+
+        credits_mask = (
+            self.prepay_credits[invoice.PREPAY_MONTH_FIELD] == self.invoice_month
+        ) & (
+            self.prepay_credits[invoice.PREPAY_GROUP_NAME_FIELD].isin(
+                managed_groups_list
+            )
+        )
+        return self.prepay_credits[credits_mask]
+
     def _backup_s3_prepay_debits(self):
         invoice_bucket = util.get_invoice_bucket()
         invoice_bucket.upload_file(
             self.prepay_debits_filepath, self.PREPAY_DEBITS_S3_BACKUP_FILEPATH
         )
 
+    def _export_prepay_credits_snapshot(self, credits_snapshot):
+        credits_snapshot.to_csv(self.CREDITS_SNAPSHOT_FILEPATH, index=False)
+
+    def _export_s3_prepay_credits_snapshot(self, credits_snapshot):
+        invoice_bucket = util.get_invoice_bucket()
+        invoice_bucket.upload_file(
+            self.CREDITS_SNAPSHOT_FILEPATH, self.CREDITS_SNAPSHOT_S3_FILEPATH
+        )
+        invoice_bucket.upload_file(
+            self.CREDITS_SNAPSHOT_FILEPATH, self.CREDITS_SNAPSHOT_S3_ARCHIVE_FILEPATH
+        )
+
     def _export_prepay_debits(self):
         self.prepay_debits.to_csv(self.prepay_debits_filepath, index=False)
 
diff --git a/process_report/tests/unit/processors/test_prepayment_processor.py b/process_report/tests/unit/processors/test_prepayment_processor.py
index 0c5ca1a..0b64e6a 100644
--- a/process_report/tests/unit/processors/test_prepayment_processor.py
+++ b/process_report/tests/unit/processors/test_prepayment_processor.py
@@ -429,3 +429,26 @@ def test_two_group_one_project(self):
             answer_prepay_debits,
             invoice_month,
         )
+
+    def test_get_credit_snapshot(self):
+        invoice_month = "2024-10"
+        test_prepay_credits = self._get_test_prepay_credits(
+            ["2024-10", "2024-10", "2024-10", "2024-09", "2024-09"],
+            ["G1", "G2", "G3", "G1", "G2"],
+            [0] * 5,
+        )
+        test_group_info_dict = {
+            "G1": {"MGHPCC Managed": True},
+            "G2": {"MGHPCC Managed": False},
+            "G3": {"MGHPCC Managed": True},
+        }
+        answer_credits_snapshot = test_prepay_credits.iloc[[0, 2]]
+
+        new_prepayment_proc = test_utils.new_prepayment_processor(
+            invoice_month=invoice_month
+        )
+        new_prepayment_proc.prepay_credits = test_prepay_credits
+        new_prepayment_proc.group_info_dict = test_group_info_dict
+        output_snapshot = new_prepayment_proc._get_prepay_credits_snapshot()
+
+        self.assertTrue(answer_credits_snapshot.equals(output_snapshot))
diff --git a/process_report/tests/util.py b/process_report/tests/util.py
index 057b9e5..b14eb39 100644
--- a/process_report/tests/util.py
+++ b/process_report/tests/util.py
@@ -168,4 +168,5 @@ def new_prepayment_processor(
         prepay_contacts,
         prepay_debits_filepath,
         upload_to_s3,
+        False,
     )