Skip to content

Commit

Permalink
Further internal refactoring of billable invoice
Browse files Browse the repository at this point in the history
  • Loading branch information
@QuanMPhm
QuanMPhm committed Jul 30, 2024
1 parent 1c4c271 commit a2f6fef
Show file tree
Hide file tree
Showing 2 changed files with 82 additions and 60 deletions.
124 changes: 70 additions & 54 deletions process_report/invoices/billable_invoice.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,20 +16,62 @@

@dataclass
class BillableInvoice(invoice.Invoice):
NEW_PI_CREDIT_CODE = "0002"
INITIAL_CREDIT_AMOUNT = 1000
EXCLUDE_SU_TYPES = ["OpenShift GPUA100SXM4", "OpenStack GPUA100SXM4"]

nonbillable_pis: list[str]
nonbillable_projects: list[str]
old_pi_filepath: str

@staticmethod
def _load_old_pis(old_pi_filepath) -> pandas.DataFrame:
try:
old_pi_df = pandas.read_csv(
old_pi_filepath,
dtype={
invoice.PI_INITIAL_CREDITS: pandas.ArrowDtype(
pyarrow.decimal128(21, 2)
),
invoice.PI_1ST_USED: pandas.ArrowDtype(pyarrow.decimal128(21, 2)),
invoice.PI_2ND_USED: pandas.ArrowDtype(pyarrow.decimal128(21, 2)),
},
)
except FileNotFoundError:
sys.exit("Applying credit 0002 failed. Old PI file does not exist")

return old_pi_df

@staticmethod
def _dump_old_pis(old_pi_filepath, old_pi_df: pandas.DataFrame):
old_pi_df.to_csv(old_pi_filepath, index=False)

def _prepare(self):
self.data = self._remove_nonbillables(
self.data, self.nonbillable_pis, self.nonbillable_projects
)
self.data = self._validate_pi_names(self.data)
self.data[invoice.CREDIT_FIELD] = None
self.data[invoice.CREDIT_CODE_FIELD] = None
self.data[invoice.BALANCE_FIELD] = Decimal(0)
self.old_pi_df = self._load_old_pis(self.old_pi_filepath)

def _process(self):
old_pi_df = self._load_old_pis(self.old_pi_filepath)
self.data, updated_old_pi_df = self._apply_credits_new_pi(self.data, old_pi_df)
self._dump_old_pis(self.old_pi_filepath, updated_old_pi_df)
self.data, self.updated_old_pi_df = self._apply_credits_new_pi(
self.data, self.old_pi_df
)

def _prepare_export(self):
self.updated_old_pi_df = self.updated_old_pi_df.astype(
{
invoice.PI_INITIAL_CREDITS: pandas.ArrowDtype(
pyarrow.decimal128(21, 2)
),
invoice.PI_1ST_USED: pandas.ArrowDtype(pyarrow.decimal128(21, 2)),
invoice.PI_2ND_USED: pandas.ArrowDtype(pyarrow.decimal128(21, 2)),
},
)
self._dump_old_pis(self.old_pi_filepath, self.updated_old_pi_df)

def _remove_nonbillables(
self,
Expand All @@ -50,47 +92,35 @@ def _validate_pi_names(self, data: pandas.DataFrame):
)
return data[~pandas.isna(data[invoice.PI_FIELD])]

def _load_old_pis(self, old_pi_filepath) -> pandas.DataFrame:
try:
old_pi_df = pandas.read_csv(
old_pi_filepath,
dtype={
invoice.PI_INITIAL_CREDITS: pandas.ArrowDtype(
pyarrow.decimal128(21, 2)
),
invoice.PI_1ST_USED: pandas.ArrowDtype(pyarrow.decimal128(21, 2)),
invoice.PI_2ND_USED: pandas.ArrowDtype(pyarrow.decimal128(21, 2)),
},
)
except FileNotFoundError:
sys.exit("Applying credit 0002 failed. Old PI file does not exist")

return old_pi_df

def _apply_credits_new_pi(
self, data: pandas.DataFrame, old_pi_df: pandas.DataFrame
):
new_pi_credit_code = "0002"
INITIAL_CREDIT_AMOUNT = 1000
EXCLUDE_SU_TYPES = ["OpenShift GPUA100SXM4", "OpenStack GPUA100SXM4"]

data[invoice.CREDIT_FIELD] = None
data[invoice.CREDIT_CODE_FIELD] = None
data[invoice.BALANCE_FIELD] = Decimal(0)

current_pi_set = set(data[invoice.PI_FIELD])
invoice_month = data[invoice.INVOICE_DATE_FIELD].iat[0]
invoice_pis = old_pi_df[old_pi_df[invoice.PI_FIRST_MONTH] == invoice_month]
if invoice_pis[invoice.PI_INITIAL_CREDITS].empty or pandas.isna(
new_pi_credit_amount := invoice_pis[invoice.PI_INITIAL_CREDITS].iat[0]
def get_initial_credit_amount(
old_pi_df, invoice_month, default_initial_credit_amount
):
new_pi_credit_amount = INITIAL_CREDIT_AMOUNT

print(f"New PI Credit set at {new_pi_credit_amount} for {invoice_month}")
first_month_processed_pis = old_pi_df[
old_pi_df[invoice.PI_FIRST_MONTH] == invoice_month
]
if first_month_processed_pis[
invoice.PI_INITIAL_CREDITS
].empty or pandas.isna(
new_pi_credit_amount := first_month_processed_pis[
invoice.PI_INITIAL_CREDITS
].iat[0]
):
new_pi_credit_amount = default_initial_credit_amount

return new_pi_credit_amount

new_pi_credit_amount = get_initial_credit_amount(
old_pi_df, self.invoice_month, self.INITIAL_CREDIT_AMOUNT
)
print(f"New PI Credit set at {new_pi_credit_amount} for {self.invoice_month}")

current_pi_set = set(data[invoice.PI_FIELD])
for pi in current_pi_set:
pi_projects = data[data[invoice.PI_FIELD] == pi]
pi_age = self._get_pi_age(old_pi_df, pi, invoice_month)
pi_age = self._get_pi_age(old_pi_df, pi, self.invoice_month)
pi_old_pi_entry = old_pi_df.loc[
old_pi_df[invoice.PI_PI_FIELD] == pi
].squeeze()
Expand All @@ -101,7 +131,7 @@ def _apply_credits_new_pi(
else:
if pi_age == 0:
if len(pi_old_pi_entry) == 0:
pi_entry = [pi, invoice_month, new_pi_credit_amount, 0, 0]
pi_entry = [pi, self.invoice_month, new_pi_credit_amount, 0, 0]
old_pi_df = pandas.concat(
[
pandas.DataFrame([pi_entry], columns=old_pi_df.columns),
Expand All @@ -126,15 +156,15 @@ def _apply_credits_new_pi(
for i, row in pi_projects.iterrows():
if (
remaining_credit == 0
or row[invoice.SU_TYPE_FIELD] in EXCLUDE_SU_TYPES
or row[invoice.SU_TYPE_FIELD] in self.EXCLUDE_SU_TYPES
):
data.at[i, invoice.BALANCE_FIELD] = row[invoice.COST_FIELD]
else:
project_cost = row[invoice.COST_FIELD]
applied_credit = min(project_cost, remaining_credit)

data.at[i, invoice.CREDIT_FIELD] = applied_credit
data.at[i, invoice.CREDIT_CODE_FIELD] = new_pi_credit_code
data.at[i, invoice.CREDIT_CODE_FIELD] = self.NEW_PI_CREDIT_CODE
data.at[i, invoice.BALANCE_FIELD] = (
row[invoice.COST_FIELD] - applied_credit
)
Expand All @@ -151,25 +181,11 @@ def _apply_credits_new_pi(
old_pi_df[invoice.PI_PI_FIELD] == pi, credit_used_field
] = credits_used

old_pi_df = old_pi_df.astype(
{
invoice.PI_INITIAL_CREDITS: pandas.ArrowDtype(
pyarrow.decimal128(21, 2)
),
invoice.PI_1ST_USED: pandas.ArrowDtype(pyarrow.decimal128(21, 2)),
invoice.PI_2ND_USED: pandas.ArrowDtype(pyarrow.decimal128(21, 2)),
},
)

return (data, old_pi_df)

def _dump_old_pis(self, old_pi_filepath, old_pi_df: pandas.DataFrame):
old_pi_df.to_csv(old_pi_filepath, index=False)

def _get_pi_age(self, old_pi_df: pandas.DataFrame, pi, invoice_month):
"""Returns time difference between current invoice month and PI's first invoice month
I.e 0 for new PIs
Will raise an error if the PI'a age is negative, which suggests a faulty invoice, or a program bug"""
first_invoice_month = old_pi_df.loc[
old_pi_df[invoice.PI_PI_FIELD] == pi, invoice.PI_FIRST_MONTH
Expand Down
18 changes: 12 additions & 6 deletions process_report/tests/unit_tests.py
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,10 @@
from unittest import TestCase, mock
import tempfile
import pandas
import pyarrow
import os
import uuid
import math
from decimal import Decimal
from textwrap import dedent

from process_report import process_report, util
Expand Down Expand Up @@ -416,6 +416,9 @@ def setUp(self):
"Balance": [10, 100, 10000, 400, 100, 0, 0, 0, 0, 200, 700],
}
self.dataframe = pandas.DataFrame(data)
self.dataframe["Credit"] = None
self.dataframe["Credit Code"] = None
self.dataframe["Balance"] = Decimal(0)
self.answer_dataframe = pandas.DataFrame(answer_df_dict)
old_pi = [
"PI,First Invoice Month,Initial Credits,1st Month Used,2nd Month Used",
Expand Down Expand Up @@ -473,9 +476,9 @@ def setUp(self):
)
.astype(
{
"Initial Credits": pandas.ArrowDtype(pyarrow.decimal128(21, 2)),
"1st Month Used": pandas.ArrowDtype(pyarrow.decimal128(21, 2)),
"2nd Month Used": pandas.ArrowDtype(pyarrow.decimal128(21, 2)),
"Initial Credits": object,
"1st Month Used": object,
"2nd Month Used": object,
},
)
.sort_values(by="PI", ignore_index=True)
Expand Down Expand Up @@ -510,6 +513,9 @@ def setUp(self):
"Cost": [500, 100, 100, 500, 500],
}
)
self.dataframe_no_gpu["Credit"] = None
self.dataframe_no_gpu["Credit Code"] = None
self.dataframe_no_gpu["Balance"] = Decimal(0)
old_pi_no_gpu = [
"PI,First Invoice Month,Initial Credits,1st Month Used,2nd Month Used",
"OldPI,2024-03,500,200,0",
Expand Down Expand Up @@ -549,7 +555,7 @@ def tearDown(self):
os.remove(self.old_pi_no_gpu_file)

def test_apply_credit_0002(self):
test_invoice = test_utils.new_billable_invoice()
test_invoice = test_utils.new_billable_invoice(invoice_month="2024-03")
old_pi_df = test_invoice._load_old_pis(self.old_pi_file)
dataframe, updated_old_pi_df = test_invoice._apply_credits_new_pi(
self.dataframe, old_pi_df
Expand All @@ -560,7 +566,7 @@ def test_apply_credit_0002(self):
self.assertTrue(self.old_pi_df_answer.equals(updated_old_pi_df))

def test_no_gpu(self):
test_invoice = test_utils.new_billable_invoice()
test_invoice = test_utils.new_billable_invoice(invoice_month="2024-03")
old_pi_df = test_invoice._load_old_pis(self.old_pi_no_gpu_file)
dataframe, _ = test_invoice._apply_credits_new_pi(
self.dataframe_no_gpu, old_pi_df
Expand Down

0 comments on commit a2f6fef

Please sign in to comment.