Add unit tests

tests/factories.py

@@ -269,6 +269,18 @@ class TaskOrderFactory(Base):
     number = factory.LazyFunction(random_task_order_number)
 
 
+class CLINFactory(Base):
+    class Meta:
+        model = CLIN
+
+    task_order = factory.SubFactory(TaskOrderFactory)
+    number = factory.LazyFunction(random_task_order_number)
+    start_date = datetime.date.today()
+    end_date = factory.LazyFunction(random_future_date)
+    obligated_amount = random.randint(100, 999999)
+    jedi_clin_type = random.choice([e.value for e in clin.JEDICLINType])
+
+
 class NotificationRecipientFactory(Base):
     class Meta:
         model = NotificationRecipient

tests/models/test_task_order.py

@@ -1,10 +1,16 @@
 from werkzeug.datastructures import FileStorage
 import pytest, datetime
 
-from atst.models.attachment import Attachment
+from atst.models import *
+from atst.models.clin import JEDICLINType
 from atst.models.task_order import TaskOrder, Status
 
-from tests.factories import random_future_date, random_past_date
+from tests.factories import (
+    CLINFactory,
+    random_future_date,
+    random_past_date,
+    TaskOrderFactory,
+)
 from tests.mocks import PDF_FILENAME
 
 
@@ -30,6 +36,33 @@ class TestTaskOrderStatus:
         assert to.status == Status.EXPIRED
 
 
+class TestBudget:
+    def test_total_contract_amount(self):
+        to = TaskOrder()
+        assert to.total_contract_amount == 0
+
+        clin1 = CLINFactory(task_order=to, jedi_clin_type=JEDICLINType.JEDI_CLIN_1)
+        clin2 = CLINFactory(task_order=to, jedi_clin_type=JEDICLINType.JEDI_CLIN_2)
+        clin3 = CLINFactory(task_order=to, jedi_clin_type=JEDICLINType.JEDI_CLIN_3)
+
+        assert (
+            to.total_contract_amount
+            == clin1.obligated_amount + clin2.obligated_amount + clin3.obligated_amount
+        )
+
+    def test_total_obligated_funds(self):
+        to = TaskOrder()
+        clin4 = CLINFactory(task_order=to, jedi_clin_type=JEDICLINType.JEDI_CLIN_4)
+        assert to.total_obligated_funds == 0
+
+        clin1 = CLINFactory(task_order=to, jedi_clin_type=JEDICLINType.JEDI_CLIN_1)
+        clin2 = CLINFactory(task_order=to, jedi_clin_type=JEDICLINType.JEDI_CLIN_2)
+        clin3 = CLINFactory(task_order=to, jedi_clin_type=JEDICLINType.JEDI_CLIN_3)
+        assert (
+            to.total_obligated_funds == clin1.obligated_amount + clin3.obligated_amount
+        )
+
+
 class TestPDF:
     def test_setting_pdf_with_attachment(self):
         to = TaskOrder()