diff --git a/alembic/versions/14cd800904bc_add_task_order_association_to_.py b/alembic/versions/14cd800904bc_add_task_order_association_to_.py new file mode 100644 index 00000000..4097c7e5 --- /dev/null +++ b/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 ### diff --git a/alembic/versions/d7db8fd35b41_add_attachment_table.py b/alembic/versions/d7db8fd35b41_add_attachment_table.py new file mode 100644 index 00000000..c3d80a48 --- /dev/null +++ b/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 ### diff --git a/atst/models/__init__.py b/atst/models/__init__.py index 9accc290..95da119b 100644 --- a/atst/models/__init__.py +++ b/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 diff --git a/atst/models/attachment.py b/atst/models/attachment.py new file mode 100644 index 00000000..21d12af0 --- /dev/null +++ b/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 diff --git a/atst/models/task_order.py b/atst/models/task_order.py index 7a9dca65..ce84a56e 100644 --- a/atst/models/task_order.py +++ b/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") diff --git a/tests/conftest.py b/tests/conftest.py index ecc12019..8093f926 100644 --- a/tests/conftest.py +++ b/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") diff --git a/tests/mocks.py b/tests/mocks.py index 2a2f0fab..63a13bd6 100644 --- a/tests/mocks.py +++ b/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" + diff --git a/tests/models/test_attachment.py b/tests/models/test_attachment.py new file mode 100644 index 00000000..e1690aa4 --- /dev/null +++ b/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) diff --git a/tests/test_uploader.py b/tests/test_uploader.py index 799cefb3..cbe39c45 100644 --- a/tests/test_uploader.py +++ b/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):