From f580f69d350c695da08d70930259107f77e39520 Mon Sep 17 00:00:00 2001
From: Patrick Smith <patrick+dds@promptworks.com>
Date: Tue, 22 Jan 2019 11:47:54 -0500
Subject: [PATCH] Raise error if csp estimate is not of proper type

---
 atst/models/task_order.py                     |  2 ++
 tests/factories.py                            | 10 +++++++
 tests/models/test_task_order.py               | 28 +++++++++++++++++++
 .../routes/task_orders/test_new_task_order.py |  3 +-
 4 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/atst/models/task_order.py b/atst/models/task_order.py
index f16b9389..ac860129 100644
--- a/atst/models/task_order.py
+++ b/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):
diff --git a/tests/factories.py b/tests/factories.py
index b456ee1d..f18e2332 100644
--- a/tests/factories.py
+++ b/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")
diff --git a/tests/models/test_task_order.py b/tests/models/test_task_order.py
index 9c0adf40..6cbfc082 100644
--- a/tests/models/test_task_order.py
+++ b/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"
diff --git a/tests/routes/task_orders/test_new_task_order.py b/tests/routes/task_orders/test_new_task_order.py
index a576611e..16a14c63 100644
--- a/tests/routes/task_orders/test_new_task_order.py
+++ b/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
     )