add attachment model and task order relation to it

alembic/versions/14cd800904bc_add_task_order_association_to_.py

@@ -0,0 +1,30 @@
+"""add task order association to attachments
+
+Revision ID: 14cd800904bc
+Revises: d7db8fd35b41
+Create Date: 2018-08-24 11:28:30.894412
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '14cd800904bc'
+down_revision = 'd7db8fd35b41'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column('task_order', sa.Column('attachment_id', sa.Integer(), nullable=True))
+    op.create_foreign_key(None, 'task_order', 'attachments', ['attachment_id'], ['id'])
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_constraint(None, 'task_order', type_='foreignkey')
+    op.drop_column('task_order', 'attachment_id')
+    # ### end Alembic commands ###

alembic/versions/d7db8fd35b41_add_attachment_table.py

@@ -0,0 +1,36 @@
+"""add attachment table
+
+Revision ID: d7db8fd35b41
+Revises: 0845b2f0f401
+Create Date: 2018-08-24 11:27:15.317181
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = 'd7db8fd35b41'
+down_revision = '0845b2f0f401'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('attachments',
+    sa.Column('id', sa.Integer(), nullable=False),
+    sa.Column('filename', sa.String(), nullable=True),
+    sa.Column('object_name', sa.String(), nullable=True),
+    sa.PrimaryKeyConstraint('id'),
+    sa.UniqueConstraint('object_name')
+    )
+    op.create_unique_constraint(None, 'task_order', ['number'])
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_constraint(None, 'task_order', type_='unique')
+    op.drop_table('attachments')
+    # ### end Alembic commands ###

atst/models/__init__.py

@@ -13,3 +13,4 @@ from .task_order import TaskOrder
 from .workspace import Workspace
 from .project import Project
 from .environment import Environment
+from .attachment import Attachment

atst/models/attachment.py

@@ -0,0 +1,32 @@
+from sqlalchemy import Column, Integer, String
+from flask import current_app as app
+
+from atst.models import Base
+from atst.database import db
+from atst.uploader import UploadError
+
+
+class AttachmentError(Exception):
+    pass
+
+
+class Attachment(Base):
+    __tablename__ = "attachments"
+
+    id = Column(Integer, primary_key=True)
+    filename = Column(String)
+    object_name = Column(String, unique=True)
+
+    @classmethod
+    def attach(cls, fyle):
+        try:
+            filename, object_name = app.uploader.upload(fyle)
+        except UploadError as e:
+            raise AttachmentError("Could not add attachment. " + str(e))
+
+        attachment = Attachment(filename=filename, object_name=object_name)
+
+        db.session.add(attachment)
+        db.session.commit()
+
+        return attachment

atst/models/task_order.py

@@ -1,6 +1,7 @@
 from enum import Enum
 
-from sqlalchemy import Column, Integer, String, Enum as SQLAEnum
+from sqlalchemy import Column, Integer, String, ForeignKey, Enum as SQLAEnum
+from sqlalchemy.orm import relationship
 
 from atst.models import Base
 
@@ -31,3 +32,6 @@ class TaskOrder(Base):
     clin_1003 = Column(Integer)
     clin_2001 = Column(Integer)
     clin_2003 = Column(Integer)
+
+    attachment_id = Column(ForeignKey("attachments.id"))
+    pdf = relationship("Attachment")

tests/conftest.py

@@ -3,25 +3,34 @@ import pytest
 import alembic.config
 import alembic.command
 from logging.config import dictConfig
+from werkzeug.datastructures import FileStorage
+from tempfile import TemporaryDirectory
 
 from atst.app import make_app, make_config
 from atst.database import db as _db
 import tests.factories as factories
+from tests.mocks import PDF_FILENAME
 
 dictConfig({"version": 1, "handlers": {"wsgi": {"class": "logging.NullHandler"}}})
 
 
 @pytest.fixture(scope="session")
 def app(request):
+    upload_dir = TemporaryDirectory()
+
     config = make_config()
+    config.update({"STORAGE_CONTAINER": upload_dir.name})
 
     _app = make_app(config)
 
     ctx = _app.app_context()
     ctx.push()
 
+
     yield _app
 
+    upload_dir.cleanup()
+
     ctx.pop()
 
 
@@ -103,3 +112,9 @@ def user_session(monkeypatch, session):
         )
 
     return set_user_session
+
+
+@pytest.fixture
+def pdf_upload():
+    with open(PDF_FILENAME, "rb") as fp:
+        yield FileStorage(fp, content_type="application/pdf")

tests/mocks.py

@@ -11,3 +11,6 @@ DOD_SDN = f"CN={DOD_SDN_INFO['last_name']}.{DOD_SDN_INFO['first_name']}.G.{DOD_S
 MOCK_VALID_PE_ID = "8675309U"
 
 FIXTURE_EMAIL_ADDRESS = "artgarfunkel@uso.mil"
+
+PDF_FILENAME = "tests/fixtures/sample.pdf"
+

tests/models/test_attachment.py

@@ -0,0 +1,18 @@
+import pytest
+from werkzeug.datastructures import FileStorage
+
+from atst.models.attachment import Attachment, AttachmentError
+
+from tests.mocks import PDF_FILENAME
+
+
+def test_attach(pdf_upload):
+    attachment = Attachment.attach(pdf_upload)
+    assert attachment.filename == PDF_FILENAME
+
+
+def test_attach_raises():
+    with open(PDF_FILENAME, "rb") as fp:
+        fs = FileStorage(fp, content_type="something/else")
+        with pytest.raises(AttachmentError):
+            Attachment.attach(fs)

tests/test_uploader.py

@@ -4,6 +4,8 @@ from werkzeug.datastructures import FileStorage
 
 from atst.uploader import Uploader, UploadError
 
+from tests.mocks import PDF_FILENAME
+
 
 @pytest.fixture(scope="function")
 def upload_dir(tmpdir):
@@ -15,22 +17,16 @@ def uploader(upload_dir):
     return Uploader("LOCAL", container=upload_dir)
 
 
-PDF_FILENAME = "tests/fixtures/sample.pdf"
 NONPDF_FILENAME = "tests/fixtures/disa-pki.html"
 
-@pytest.fixture
-def pdf():
-    with open(PDF_FILENAME, "rb") as fp:
-        yield FileStorage(fp, content_type="application/pdf")
 
-
-def test_upload(uploader, upload_dir, pdf):
-    filename, object_name = uploader.upload(pdf)
+def test_upload(uploader, upload_dir, pdf_upload):
+    filename, object_name = uploader.upload(pdf_upload)
     assert filename == PDF_FILENAME
     assert os.path.isfile(os.path.join(upload_dir, object_name))
 
 
-def test_upload_fails_for_non_pdfs(uploader, pdf):
+def test_upload_fails_for_non_pdfs(uploader):
     with open(NONPDF_FILENAME, "rb") as fp:
         fs = FileStorage(fp, content_type="text/plain")
         with pytest.raises(UploadError):