diff --git a/atst/domain/csp/file_uploads.py b/atst/domain/csp/file_uploads.py index d9c2a367..57ea8f0e 100644 --- a/atst/domain/csp/file_uploads.py +++ b/atst/domain/csp/file_uploads.py @@ -7,10 +7,13 @@ import boto3 def build_uploader(config): - if config["CSP"] == "aws": + csp = config.get("CSP") + if csp == "aws": return AwsUploader(config) - elif config["CSP"] == "azure": + elif csp == "azure": return AzureUploader(config) + else: + return MockUploader(config) class Uploader: @@ -21,6 +24,14 @@ class Uploader: return str(uuid4()) +class MockUploader(Uploader): + def __init__(self, config): + self.config = config + + def get_token(self): + return ({}, self.object_name()) + + class AzureUploader(Uploader): def __init__(self, config): self.config = config diff --git a/atst/forms/task_order.py b/atst/forms/task_order.py index 3ab8684c..c8fcbdb9 100644 --- a/atst/forms/task_order.py +++ b/atst/forms/task_order.py @@ -2,20 +2,17 @@ from wtforms.fields import ( BooleanField, DecimalField, FieldList, - FileField, FormField, StringField, - HiddenField + HiddenField, ) from wtforms.fields.html5 import DateField from wtforms.validators import Required, Optional -from flask_wtf.file import FileAllowed from flask_wtf import FlaskForm from .data import JEDI_CLIN_TYPES from .fields import SelectField from .forms import BaseForm -from atst.forms.validators import FileLength from atst.utils.localization import translate @@ -70,7 +67,11 @@ class AttachmentForm(BaseForm): class TaskOrderForm(BaseForm): number = StringField(label=translate("forms.task_order.number_description")) - pdf = FormField(AttachmentForm, label=translate("task_orders.form.supporting_docs_size_limit"), description=translate("task_orders.form.supporting_docs_size_limit")) + pdf = FormField( + AttachmentForm, + label=translate("task_orders.form.supporting_docs_size_limit"), + description=translate("task_orders.form.supporting_docs_size_limit"), + ) clins = FieldList(FormField(CLINForm)) diff --git a/atst/models/attachment.py b/atst/models/attachment.py index e4a9d6c2..4edb3441 100644 --- a/atst/models/attachment.py +++ b/atst/models/attachment.py @@ -40,6 +40,16 @@ class Attachment(Base, mixins.TimestampsMixin): return attachment + @classmethod + def get_or_create(cls, object_name, params): + try: + return db.session.query(Attachment).filter_by(object_name=object_name).one() + except NoResultFound: + new_attachment = cls(**params) + db.session.add(new_attachment) + db.session.commit() + return new_attachment + @classmethod def get(cls, id_): try: diff --git a/atst/models/task_order.py b/atst/models/task_order.py index 5b873d71..8b060164 100644 --- a/atst/models/task_order.py +++ b/atst/models/task_order.py @@ -9,7 +9,6 @@ from werkzeug.datastructures import FileStorage from atst.models import Attachment, Base, mixins, types from atst.models.clin import JEDICLINType from atst.utils.clock import Clock -from atst.database import db class Status(Enum): @@ -56,15 +55,18 @@ class TaskOrder(Base, mixins.TimestampsMixin): self._pdf = self._set_attachment(new_pdf, "_pdf") def _set_attachment(self, new_attachment, attribute): - if isinstance(new_attachment, dict): - attachment = Attachment(**new_attachment) - db.session.add(attachment) - db.session.commit - return attachment if isinstance(new_attachment, Attachment): return new_attachment elif isinstance(new_attachment, FileStorage): return Attachment.attach(new_attachment, "task_order", self.id) + elif isinstance(new_attachment, dict): + if new_attachment["filename"] and new_attachment["object_name"]: + attachment = Attachment.get_or_create( + new_attachment["object_name"], new_attachment + ) + return attachment + else: + return None elif not new_attachment and hasattr(self, attribute): return None else: diff --git a/atst/routes/task_orders/new.py b/atst/routes/task_orders/new.py index 375fc845..f8a9a3a9 100644 --- a/atst/routes/task_orders/new.py +++ b/atst/routes/task_orders/new.py @@ -6,9 +6,6 @@ from flask import ( url_for, current_app, ) -from azure.storage.common import CloudStorageAccount -from azure.storage.blob import ContainerPermissions -from datetime import datetime, timedelta from . import task_orders_bp from atst.domain.authz.decorator import user_can_access_decorator as user_can @@ -100,7 +97,6 @@ def edit(task_order_id): @task_orders_bp.route("/task_orders//form/step_1") @user_can(Permissions.CREATE_TASK_ORDER, message="view task order form") def form_step_one_add_pdf(portfolio_id=None, task_order_id=None): - return render_task_orders_edit( "task_orders/step_1.html", portfolio_id=portfolio_id, diff --git a/tests/factories.py b/tests/factories.py index 4dd2a7f8..91fe0566 100644 --- a/tests/factories.py +++ b/tests/factories.py @@ -269,7 +269,9 @@ class TaskOrderFactory(Base): class Meta: model = TaskOrder - portfolio = factory.SubFactory(PortfolioFactory) + portfolio = factory.SubFactory( + PortfolioFactory, owner=factory.SelfAttribute("..creator") + ) number = factory.LazyFunction(random_task_order_number) creator = factory.SubFactory(UserFactory) _pdf = factory.SubFactory(AttachmentFactory) diff --git a/tests/routes/task_orders/test_new.py b/tests/routes/task_orders/test_new.py index b88bbf7b..9cb5a6b4 100644 --- a/tests/routes/task_orders/test_new.py +++ b/tests/routes/task_orders/test_new.py @@ -17,6 +17,10 @@ from tests.factories import ( ) +def build_pdf_form_data(filename="sample.pdf", object_name="object_name"): + return {"pdf-filename": filename, "pdf-object_name": object_name} + + @pytest.fixture def task_order(): user = UserFactory.create() @@ -55,19 +59,16 @@ def test_task_orders_form_step_one_add_pdf(client, user_session, portfolio): assert response.status_code == 200 -def test_task_orders_submit_form_step_one_add_pdf( - client, user_session, portfolio, pdf_upload, session -): +def test_task_orders_submit_form_step_one_add_pdf(client, user_session, portfolio): user_session(portfolio.owner) - form_data = {"pdf": pdf_upload} response = client.post( url_for("task_orders.submit_form_step_one_add_pdf", portfolio_id=portfolio.id), - data=form_data, + data=build_pdf_form_data(), ) assert response.status_code == 302 task_order = portfolio.task_orders[0] - assert task_order.pdf.filename == pdf_upload.filename + assert task_order.pdf.filename == "sample.pdf" def test_task_orders_form_step_one_add_pdf_existing_to( @@ -80,35 +81,29 @@ def test_task_orders_form_step_one_add_pdf_existing_to( assert response.status_code == 200 -def test_task_orders_submit_form_step_one_add_pdf_existing_to( - client, user_session, task_order, pdf_upload, pdf_upload2 -): - task_order.pdf = pdf_upload - assert task_order.pdf.filename == pdf_upload.filename - +def test_task_orders_submit_form_step_one_add_pdf_existing_to(client, user_session): + task_order = TaskOrderFactory.create() user_session(task_order.creator) - form_data = {"pdf": pdf_upload2} response = client.post( url_for( "task_orders.submit_form_step_one_add_pdf", task_order_id=task_order.id ), - data=form_data, + data=build_pdf_form_data(), ) assert response.status_code == 302 - assert task_order.pdf.filename == pdf_upload2.filename + assert task_order.pdf.filename == "sample.pdf" def test_task_orders_submit_form_step_one_add_pdf_delete_pdf( client, user_session, portfolio, pdf_upload ): user_session(portfolio.owner) - task_order = TaskOrderFactory.create(pdf=pdf_upload, portfolio=portfolio) - data = {"pdf": ""} + task_order = TaskOrderFactory.create(portfolio=portfolio) response = client.post( url_for( "task_orders.submit_form_step_one_add_pdf", task_order_id=task_order.id ), - data=data, + data=build_pdf_form_data(filename="", object_name=""), ) assert task_order.pdf is None assert response.status_code == 302 diff --git a/tests/test_access.py b/tests/test_access.py index eefa2160..a47c7b02 100644 --- a/tests/test_access.py +++ b/tests/test_access.py @@ -438,9 +438,7 @@ def test_task_orders_download_task_order_pdf_access(get_url_assert_status, monke rando = user_with() portfolio = PortfolioFactory.create(owner=owner) - task_order = TaskOrderFactory.create( - portfolio=portfolio, pdf=AttachmentFactory.create() - ) + task_order = TaskOrderFactory.create(portfolio=portfolio) url = url_for("task_orders.download_task_order_pdf", task_order_id=task_order.id) get_url_assert_status(owner, url, 200)