diff --git a/tests/factories.py b/tests/factories.py index 341887c0..e91ec26b 100644 --- a/tests/factories.py +++ b/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 diff --git a/tests/models/test_task_order.py b/tests/models/test_task_order.py index 3a25781a..d6ae9e1c 100644 --- a/tests/models/test_task_order.py +++ b/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()