diff --git a/atst/models/task_order.py b/atst/models/task_order.py index 61bd4b60..f16b9389 100644 --- a/atst/models/task_order.py +++ b/atst/models/task_order.py @@ -5,6 +5,7 @@ from sqlalchemy import Column, Numeric, String, ForeignKey, Date, Integer from sqlalchemy.ext.hybrid import hybrid_property from sqlalchemy.types import ARRAY from sqlalchemy.orm import relationship +from werkzeug.datastructures import FileStorage from atst.models import Attachment, Base, types, mixins @@ -81,7 +82,7 @@ class TaskOrder(Base, mixins.TimestampsMixin): def csp_estimate(self, new_csp_estimate): if isinstance(new_csp_estimate, Attachment): self._csp_estimate = new_csp_estimate - else: + elif isinstance(new_csp_estimate, FileStorage): self._csp_estimate = Attachment.attach( new_csp_estimate, "task_order", self.id ) diff --git a/tests/domain/test_task_orders.py b/tests/domain/test_task_orders.py index d992050b..ab92e024 100644 --- a/tests/domain/test_task_orders.py +++ b/tests/domain/test_task_orders.py @@ -2,6 +2,7 @@ import pytest from atst.domain.task_orders import TaskOrders, TaskOrderError from atst.domain.exceptions import UnauthorizedError +from atst.models.attachment import Attachment from tests.factories import ( TaskOrderFactory, @@ -26,10 +27,18 @@ def test_is_section_complete(): def test_all_sections_complete(): task_order = TaskOrderFactory.create() + attachment = Attachment( + filename="sample_attachment", + object_name="sample", + resource="task_order", + resource_id=task_order.id, + ) + + custom_attrs = {"csp_estimate": attachment} for attr_list in TaskOrders.SECTIONS.values(): for attr in attr_list: if not getattr(task_order, attr): - setattr(task_order, attr, "str12345") + setattr(task_order, attr, custom_attrs.get(attr, "str12345")) task_order.scope = None assert not TaskOrders.all_sections_complete(task_order) diff --git a/tests/routes/task_orders/test_new_task_order.py b/tests/routes/task_orders/test_new_task_order.py index 53374ed6..a576611e 100644 --- a/tests/routes/task_orders/test_new_task_order.py +++ b/tests/routes/task_orders/test_new_task_order.py @@ -2,6 +2,7 @@ import pytest from flask import url_for from atst.domain.task_orders import TaskOrders +from atst.models.attachment import Attachment from atst.routes.task_orders.new import ShowTaskOrderWorkflow, UpdateTaskOrderWorkflow from tests.factories import UserFactory, TaskOrderFactory, PortfolioFactory @@ -124,8 +125,11 @@ def test_task_order_form_shows_errors(client, user_session): def task_order(): user = UserFactory.create() portfolio = PortfolioFactory.create(owner=user) + attachment = Attachment(filename="sample_attachment", object_name="sample") - return TaskOrderFactory.create(creator=user, portfolio=portfolio) + return TaskOrderFactory.create( + creator=user, portfolio=portfolio, csp_estimate=attachment + ) def test_show_task_order(task_order):