From 692d321cef98324dd78a793d0927def71c408c85 Mon Sep 17 00:00:00 2001
From: Caitlin Barnard <caitlin@madetech.com>
Date: Tue, 10 Dec 2024 11:13:25 +0000
Subject: [PATCH] Add unique constraint and refactor

---
 ...17_alter_attrition_cost_centre_and_more.py | 44 +++++++++++++++++++
 core/models.py                                |  9 ++--
 payroll/services/payroll.py                   | 14 +++---
 3 files changed, 56 insertions(+), 11 deletions(-)
 create mode 100644 core/migrations/0017_alter_attrition_cost_centre_and_more.py

diff --git a/core/migrations/0017_alter_attrition_cost_centre_and_more.py b/core/migrations/0017_alter_attrition_cost_centre_and_more.py
new file mode 100644
index 00000000..25a6d46b
--- /dev/null
+++ b/core/migrations/0017_alter_attrition_cost_centre_and_more.py
@@ -0,0 +1,44 @@
+# Generated by Django 5.1.3 on 2024-12-10 11:10
+
+import django.db.models.deletion
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("core", "0016_alter_attrition_options"),
+        ("costcentre", "0008_alter_simplehistoryarchivedcostcentre_options_and_more"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="attrition",
+            name="cost_centre",
+            field=models.ForeignKey(
+                blank=True,
+                null=True,
+                on_delete=django.db.models.deletion.PROTECT,
+                to="costcentre.costcentre",
+                unique=True,
+            ),
+        ),
+        migrations.AlterField(
+            model_name="attrition",
+            name="financial_year",
+            field=models.ForeignKey(
+                on_delete=django.db.models.deletion.PROTECT,
+                to="core.financialyear",
+                unique=True,
+            ),
+        ),
+        migrations.AlterField(
+            model_name="payuplift",
+            name="financial_year",
+            field=models.ForeignKey(
+                on_delete=django.db.models.deletion.PROTECT,
+                to="core.financialyear",
+                unique=True,
+            ),
+        ),
+    ]
diff --git a/core/models.py b/core/models.py
index a248682c..ddb438a9 100644
--- a/core/models.py
+++ b/core/models.py
@@ -98,8 +98,7 @@ def periods(self) -> list[float]:
         return [getattr(self, month) for month in MONTHS]
 
     financial_year = models.ForeignKey(
-        FinancialYear,
-        on_delete=models.PROTECT,
+        FinancialYear, on_delete=models.PROTECT, unique=True
     )
     apr = models.FloatField(default=1.0)
     may = models.FloatField(default=1.0)
@@ -124,7 +123,11 @@ class Meta:
         verbose_name_plural = "attrition"
 
     cost_centre = models.ForeignKey(
-        "costcentre.CostCentre", on_delete=models.PROTECT, null=True, blank=True
+        "costcentre.CostCentre",
+        on_delete=models.PROTECT,
+        null=True,
+        blank=True,
+        unique=True,
     )
 
 
diff --git a/payroll/services/payroll.py b/payroll/services/payroll.py
index 80605ea5..13e412e3 100644
--- a/payroll/services/payroll.py
+++ b/payroll/services/payroll.py
@@ -80,12 +80,8 @@ def payroll_forecast_report(
     pay_uplift_obj = PayUplift.objects.filter(financial_year=financial_year).first()
     attrition_obj = get_attrition_instance(financial_year, cost_centre)
 
-    pay_uplift = (
-        np.array(pay_uplift_obj.periods) if pay_uplift_obj is not None else np.ones(12)
-    )
-    attrition = (
-        np.array(attrition_obj.periods) if attrition_obj is not None else np.ones(12)
-    )
+    pay_uplift = np.array(pay_uplift_obj.periods) if pay_uplift_obj else np.ones(12)
+    attrition = np.array(attrition_obj.periods) if attrition_obj else np.ones(12)
     attrition_accumulate = np.array(list(accumulate(attrition, operator.mul)))
 
     for employee in employee_qs.iterator():
@@ -124,12 +120,14 @@ def payroll_forecast_report(
             )
 
 
-def get_attrition_instance(financial_year, cost_centre):
+def get_attrition_instance(
+    financial_year: FinancialYear, cost_centre: CostCentre
+) -> Attrition | None:
     instance = Attrition.objects.filter(
         financial_year=financial_year, cost_centre=cost_centre
     ).first()
 
-    if instance is not None:
+    if instance:
         return instance
 
     return Attrition.objects.filter(