Raise error if csp estimate is not of proper type

atst/models/task_order.py

@@ -86,6 +86,8 @@ class TaskOrder(Base, mixins.TimestampsMixin):
             self._csp_estimate = Attachment.attach(
                 new_csp_estimate, "task_order", self.id
             )
+        else:
+            raise TypeError("Could not set csp_estimate with invalid type")
 
     @property
     def is_submitted(self):

tests/factories.py

@@ -7,6 +7,7 @@ import datetime
 from faker import Faker as _Faker
 
 from atst.forms import data
+from atst.models.attachment import Attachment
 from atst.models.environment import Environment
 from atst.models.request import Request
 from atst.models.request_revision import RequestRevision
@@ -375,6 +376,14 @@ class InvitationFactory(Base):
     expiration_time = Invitations.current_expiration_time()
 
 
+class AttachmentFactory(Base):
+    class Meta:
+        model = Attachment
+
+    filename = factory.Faker("domain_word")
+    object_name = factory.Faker("domain_word")
+
+
 class TaskOrderFactory(Base):
     class Meta:
         model = TaskOrder
@@ -401,6 +410,7 @@ class TaskOrderFactory(Base):
         lambda *args: random_future_date(year_min=2, year_max=5)
     )
     performance_length = random.randint(1, 24)
+    csp_estimate = factory.SubFactory(AttachmentFactory)
 
     ko_first_name = factory.Faker("first_name")
     ko_last_name = factory.Faker("last_name")

tests/models/test_task_order.py

@@ -1,6 +1,11 @@
+from werkzeug.datastructures import FileStorage
+import pytest
+
+from atst.models.attachment import Attachment
 from atst.models.task_order import TaskOrder, Status
 
 from tests.factories import random_future_date, random_past_date
+from tests.mocks import PDF_FILENAME
 
 
 class TestTaskOrderStatus:
@@ -30,3 +35,26 @@ def test_is_submitted():
 
     to = TaskOrder(number="42")
     assert to.is_submitted
+
+
+class TestCSPEstimate:
+    def test_setting_estimate_with_attachment(self):
+        to = TaskOrder()
+        attachment = Attachment(filename="sample.pdf", object_name="sample")
+        to.csp_estimate = attachment
+
+        assert to.attachment_id == attachment.id
+
+    def test_setting_estimate_with_file_storage(self):
+        to = TaskOrder()
+        with open(PDF_FILENAME, "rb") as fp:
+            fs = FileStorage(fp, content_type="application/pdf")
+            to.csp_estimate = fs
+
+        assert to.csp_estimate is not None
+        assert to.csp_estimate.filename == PDF_FILENAME
+
+    def test_setting_estimate_with_invalid_object(self):
+        to = TaskOrder()
+        with pytest.raises(TypeError):
+            to.csp_estimate = "invalid"

tests/routes/task_orders/test_new_task_order.py

@@ -43,7 +43,7 @@ def serialize_dates(data):
 
 # TODO: this test will need to be more complicated when we add validation to
 # the forms
-def test_create_new_task_order(client, user_session):
+def test_create_new_task_order(client, user_session, pdf_upload):
     creator = UserFactory.create()
     user_session(creator)
 
@@ -66,6 +66,7 @@ def test_create_new_task_order(client, user_session):
 
     funding_data = slice_data_for_section(task_order_data, "funding")
     funding_data = serialize_dates(funding_data)
+    funding_data["csp_estimate"] = pdf_upload
     response = client.post(
         response.headers["Location"], data=funding_data, follow_redirects=False
     )