From 9182b1078c2d37691d1dc5a5cebea8a0160091cd Mon Sep 17 00:00:00 2001
From: Montana <montana-ctr@friends.dds.mil>
Date: Wed, 6 Feb 2019 08:55:46 -0500
Subject: [PATCH] Refactor attachment setters and add tests

---
 atst/models/task_order.py              | 33 ++++++++++-----------
 tests/models/test_task_order.py        | 40 +++++++++++++++++++++++++-
 tests/routes/task_orders/test_index.py |  2 +-
 3 files changed, 57 insertions(+), 18 deletions(-)

diff --git a/atst/models/task_order.py b/atst/models/task_order.py
index 60cb2ac9..119cf217 100644
--- a/atst/models/task_order.py
+++ b/atst/models/task_order.py
@@ -84,15 +84,7 @@ class TaskOrder(Base, mixins.TimestampsMixin):
 
     @csp_estimate.setter
     def csp_estimate(self, new_csp_estimate):
-        if isinstance(new_csp_estimate, Attachment):
-            self._csp_estimate = new_csp_estimate
-        elif isinstance(new_csp_estimate, FileStorage):
-            self._csp_estimate = Attachment.attach(
-                new_csp_estimate, "task_order", self.id
-            )
-        elif not new_csp_estimate and self._csp_estimate:
-            self._csp_estimate = None
-        elif new_csp_estimate:
+        if not self._set_attachment_type(new_csp_estimate, "_csp_estimate"):
             raise TypeError("Could not set csp_estimate with invalid type")
 
     @hybrid_property
@@ -101,15 +93,24 @@ class TaskOrder(Base, mixins.TimestampsMixin):
 
     @pdf.setter
     def pdf(self, new_pdf):
-        if isinstance(new_pdf, Attachment):
-            self._pdf = new_pdf
-        elif isinstance(new_pdf, FileStorage):
-            self._pdf = Attachment.attach(new_pdf, "task_order", self.id)
-        elif not new_pdf and self._pdf:
-            self._pdf = None
-        elif new_pdf:
+        if not self._set_attachment_type(new_pdf, "_pdf"):
             raise TypeError("Could not set pdf with invalid type")
 
+    def _set_attachment_type(self, new_attachment, property):
+        if isinstance(new_attachment, Attachment):
+            setattr(self, property, new_attachment)
+            return True
+        elif isinstance(new_attachment, FileStorage):
+            setattr(
+                self, property, Attachment.attach(new_attachment, "task_order", self.id)
+            )
+            return True
+        elif not new_attachment and hasattr(self, property):
+            setattr(self, property, None)
+            return True
+        else:
+            return False
+
     @property
     def is_submitted(self):
 
diff --git a/tests/models/test_task_order.py b/tests/models/test_task_order.py
index 560c6989..908daec2 100644
--- a/tests/models/test_task_order.py
+++ b/tests/models/test_task_order.py
@@ -47,7 +47,7 @@ class TestCSPEstimate:
         attachment = Attachment(filename="sample.pdf", object_name="sample")
         to.csp_estimate = attachment
 
-        assert to.attachment_id == attachment.id
+        assert to.csp_attachment_id == attachment.id
 
     def test_setting_estimate_with_file_storage(self):
         to = TaskOrder()
@@ -77,3 +77,41 @@ class TestCSPEstimate:
 
         to.csp_estimate = ""
         assert to.csp_estimate is None
+
+
+class TestPDF:
+    def test_setting_pdf_with_attachment(self):
+        to = TaskOrder()
+        attachment = Attachment(filename="sample.pdf", object_name="sample")
+        to.pdf = attachment
+
+        assert to.pdf_attachment_id == attachment.id
+
+    def test_setting_pdf_with_file_storage(self):
+        to = TaskOrder()
+        with open(PDF_FILENAME, "rb") as fp:
+            fs = FileStorage(fp, content_type="application/pdf")
+            to.pdf = fs
+
+        assert to.pdf is not None
+        assert to.pdf.filename == PDF_FILENAME
+
+    def test_setting_pdf_with_invalid_object(self):
+        to = TaskOrder()
+        with pytest.raises(TypeError):
+            to.pdf = "invalid"
+
+    def test_setting_pdf_with_empty_value(self):
+        to = TaskOrder()
+        assert to.pdf is None
+
+        to.pdf = ""
+        assert to.pdf is None
+
+    def test_removing_pdf(self):
+        attachment = Attachment(filename="sample.pdf", object_name="sample")
+        to = TaskOrder(pdf=attachment)
+        assert to.pdf is not None
+
+        to.pdf = ""
+        assert to.pdf is None
diff --git a/tests/routes/task_orders/test_index.py b/tests/routes/task_orders/test_index.py
index 45db3a82..784e4488 100644
--- a/tests/routes/task_orders/test_index.py
+++ b/tests/routes/task_orders/test_index.py
@@ -53,7 +53,7 @@ class TestDownloadCSPEstimate:
         assert expected_contents == response.data
 
     def test_download_without_attachment(self, client, user_session):
-        self.task_order.attachment_id = None
+        self.task_order.csp_attachment_id = None
         user_session(self.user)
         response = client.get(
             url_for(