diff --git a/atst/routes/task_orders/index.py b/atst/routes/task_orders/index.py index 6a632cbb..6abf34fd 100644 --- a/atst/routes/task_orders/index.py +++ b/atst/routes/task_orders/index.py @@ -3,6 +3,7 @@ from flask import g, Response, current_app as app from . import task_orders_bp from atst.domain.task_orders import TaskOrders +from atst.domain.exceptions import NotFoundError from atst.utils.docx import Docx diff --git a/tests/routes/task_orders/test_index.py b/tests/routes/task_orders/test_index.py index 9ca5ea64..45db3a82 100644 --- a/tests/routes/task_orders/test_index.py +++ b/tests/routes/task_orders/test_index.py @@ -28,3 +28,46 @@ def test_download_summary(client, user_session): for attr, val in task_order.to_dictionary().items(): assert attr in doc assert xml_translated(val) in doc + + +class TestDownloadCSPEstimate: + def setup(self): + self.user = UserFactory.create() + self.portfolio = PortfolioFactory.create(owner=self.user) + self.task_order = TaskOrderFactory.create( + creator=self.user, portfolio=self.portfolio + ) + + def test_successful_download(self, client, user_session, pdf_upload): + self.task_order.csp_estimate = pdf_upload + user_session(self.user) + response = client.get( + url_for( + "task_orders.download_csp_estimate", task_order_id=self.task_order.id + ) + ) + assert response.status_code == 200 + + pdf_upload.seek(0) + expected_contents = pdf_upload.read() + assert expected_contents == response.data + + def test_download_without_attachment(self, client, user_session): + self.task_order.attachment_id = None + user_session(self.user) + response = client.get( + url_for( + "task_orders.download_csp_estimate", task_order_id=self.task_order.id + ) + ) + assert response.status_code == 404 + + def test_download_with_wrong_user(self, client, user_session): + other_user = UserFactory.create() + user_session(other_user) + response = client.get( + url_for( + "task_orders.download_csp_estimate", task_order_id=self.task_order.id + ) + ) + assert response.status_code == 404