Merge pull request #511 from dod-ccpo/csp-integration

Refactor CSP integration
This commit is contained in:
patricksmithdds 2019-01-03 09:59:55 -05:00 committed by GitHub
commit 66c88d2869
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
11 changed files with 381 additions and 285 deletions

View File

@ -21,9 +21,9 @@ from atst.routes.errors import make_error_pages
from atst.domain.authnid.crl import CRLCache from atst.domain.authnid.crl import CRLCache
from atst.domain.auth import apply_authentication from atst.domain.auth import apply_authentication
from atst.domain.authz import Authorization from atst.domain.authz import Authorization
from atst.domain.csp import make_csp_provider
from atst.models.permissions import Permissions from atst.models.permissions import Permissions
from atst.eda_client import MockEDAClient from atst.eda_client import MockEDAClient
from atst.uploader import Uploader
from atst.utils import mailer from atst.utils import mailer
from atst.utils.form_cache import FormCache from atst.utils.form_cache import FormCache
from atst.queue import queue from atst.queue import queue
@ -52,7 +52,7 @@ def make_app(config):
make_crl_validator(app) make_crl_validator(app)
register_filters(app) register_filters(app)
make_eda_client(app) make_eda_client(app)
make_upload_storage(app) make_csp_provider(app)
make_mailer(app) make_mailer(app)
queue.init_app(app) queue.init_app(app)
@ -191,16 +191,6 @@ def make_eda_client(app):
app.eda_client = MockEDAClient() app.eda_client = MockEDAClient()
def make_upload_storage(app):
uploader = Uploader(
provider=app.config.get("STORAGE_PROVIDER"),
container=app.config.get("STORAGE_CONTAINER"),
key=app.config.get("STORAGE_KEY"),
secret=app.config.get("STORAGE_SECRET"),
)
app.uploader = uploader
def make_mailer(app): def make_mailer(app):
if app.config["DEBUG"]: if app.config["DEBUG"]:
mailer_connection = mailer.RedisConnection(app.redis) mailer_connection = mailer.RedisConnection(app.redis)

View File

@ -0,0 +1,12 @@
from .files import RackspaceFileProvider
from .reports import MockReportingProvider
class MockCSP:
def __init__(self, app):
self.files = RackspaceFileProvider(app)
self.reports = MockReportingProvider()
def make_csp_provider(app):
app.csp = MockCSP(app)

View File

@ -4,18 +4,13 @@ from uuid import uuid4
from libcloud.storage.types import Provider from libcloud.storage.types import Provider
from libcloud.storage.providers import get_driver from libcloud.storage.providers import get_driver
from atst.domain.exceptions import UploadError
class UploadError(Exception):
pass
class Uploader: class FileProviderInterface:
_PERMITTED_MIMETYPES = ["application/pdf"] _PERMITTED_MIMETYPES = ["application/pdf"]
def __init__(self, provider, container=None, key=None, secret=None): def _enforce_mimetype(self, fyle):
self.container = self._get_container(provider, container, key, secret)
def upload(self, fyle):
# TODO: for hardening, we should probably use a better library for # TODO: for hardening, we should probably use a better library for
# determining mimetype and not rely on FileUpload's determination # determining mimetype and not rely on FileUpload's determination
# TODO: we should set MAX_CONTENT_LENGTH in the config to prevent large # TODO: we should set MAX_CONTENT_LENGTH in the config to prevent large
@ -27,6 +22,38 @@ class Uploader:
) )
) )
def upload(self, fyle): # pragma: no cover
"""Store the file object `fyle` in the CSP. This method returns the
object name that can be used to later look up the file."""
raise NotImplementedError()
def download(self, object_name): # pragma: no cover
"""Retrieve the stored file represented by `object_name`. Returns a
file object.
"""
raise NotImplementedError()
class RackspaceFileProvider(FileProviderInterface):
def __init__(self, app):
self.container = self._get_container(
provider=app.config.get("STORAGE_PROVIDER"),
container=app.config.get("STORAGE_CONTAINER"),
key=app.config.get("STORAGE_KEY"),
secret=app.config.get("STORAGE_SECRET"),
)
def _get_container(self, provider, container=None, key=None, secret=None):
if provider == "LOCAL": # pragma: no branch
key = container
container = ""
driver = get_driver(getattr(Provider, provider))(key=key, secret=secret)
return driver.get_container(container)
def upload(self, fyle):
self._enforce_mimetype(fyle)
object_name = uuid4().hex object_name = uuid4().hex
with NamedTemporaryFile() as tempfile: with NamedTemporaryFile() as tempfile:
tempfile.write(fyle.stream.read()) tempfile.write(fyle.stream.read())
@ -35,18 +62,10 @@ class Uploader:
object_name=object_name, object_name=object_name,
extra={"acl": "private"}, extra={"acl": "private"},
) )
return (fyle.filename, object_name) return object_name
def download_stream(self, object_name): def download(self, object_name):
obj = self.container.get_object(object_name=object_name) obj = self.container.get_object(object_name=object_name)
with NamedTemporaryFile() as tempfile: with NamedTemporaryFile() as tempfile:
obj.download(tempfile.name, overwrite_existing=True) obj.download(tempfile.name, overwrite_existing=True)
return open(tempfile.name, "rb") return open(tempfile.name, "rb")
def _get_container(self, provider, container, key, secret):
if provider == "LOCAL":
key = container
container = ""
driver = get_driver(getattr(Provider, provider))(key=key, secret=secret)
return driver.get_container(container)

305
atst/domain/csp/reports.py Normal file
View File

@ -0,0 +1,305 @@
from itertools import groupby
from collections import OrderedDict
import pendulum
class ReportingInterface:
def monthly_totals_for_environment(environment):
"""Return the monthly totals for the specified environment.
Data should be in the format of a dictionary with the month as the key
and the spend in that month as the value. For example:
{ "01/2018": 79.85, "02/2018": 86.54 }
"""
raise NotImplementedError()
class MockEnvironment:
def __init__(self, id_, env_name):
self.id = id_
self.name = env_name
class MockProject:
def __init__(self, project_name, envs):
def make_env(name):
return MockEnvironment("{}_{}".format(project_name, name), name)
self.name = project_name
self.environments = [make_env(env_name) for env_name in envs]
class MockReportingProvider(ReportingInterface):
MONTHLY_SPEND_BY_ENVIRONMENT = {
"LC04_Integ": {
"02/2018": 284,
"03/2018": 1210,
"04/2018": 1430,
"05/2018": 1366,
"06/2018": 1169,
"07/2018": 991,
"08/2018": 978,
"09/2018": 737,
},
"LC04_PreProd": {
"02/2018": 812,
"03/2018": 1389,
"04/2018": 1425,
"05/2018": 1306,
"06/2018": 1112,
"07/2018": 936,
"08/2018": 921,
"09/2018": 694,
},
"LC04_Prod": {
"02/2018": 1742,
"03/2018": 1716,
"04/2018": 1866,
"05/2018": 1809,
"06/2018": 1839,
"07/2018": 1633,
"08/2018": 1654,
"09/2018": 1103,
},
"SF18_Integ": {
"04/2018": 1498,
"05/2018": 1400,
"06/2018": 1394,
"07/2018": 1171,
"08/2018": 1200,
"09/2018": 963,
},
"SF18_PreProd": {
"04/2018": 1780,
"05/2018": 1667,
"06/2018": 1703,
"07/2018": 1474,
"08/2018": 1441,
"09/2018": 933,
},
"SF18_Prod": {
"04/2018": 1686,
"05/2018": 1779,
"06/2018": 1792,
"07/2018": 1570,
"08/2018": 1539,
"09/2018": 986,
},
"Canton_Prod": {
"05/2018": 28699,
"06/2018": 26766,
"07/2018": 22619,
"08/2018": 24090,
"09/2018": 16719,
},
"BD04_Integ": {},
"BD04_PreProd": {
"02/2018": 7019,
"03/2018": 3004,
"04/2018": 2691,
"05/2018": 2901,
"06/2018": 3463,
"07/2018": 3314,
"08/2018": 3432,
"09/2018": 723,
},
"SCV18_Dev": {"05/2019": 9797},
"Crown_CR Portal Dev": {
"03/2018": 208,
"04/2018": 457,
"05/2018": 671,
"06/2018": 136,
"07/2018": 1524,
"08/2018": 2077,
"09/2018": 1858,
},
"Crown_CR Staging": {
"03/2018": 208,
"04/2018": 457,
"05/2018": 671,
"06/2018": 136,
"07/2018": 1524,
"08/2018": 2077,
"09/2018": 1858,
},
"Crown_CR Portal Test 1": {"07/2018": 806, "08/2018": 1966, "09/2018": 2597},
"Crown_Jewels Prod": {"07/2018": 806, "08/2018": 1966, "09/2018": 2597},
"Crown_Jewels Dev": {
"03/2018": 145,
"04/2018": 719,
"05/2018": 1243,
"06/2018": 2214,
"07/2018": 2959,
"08/2018": 4151,
"09/2018": 4260,
},
"NP02_Integ": {"08/2018": 284, "09/2018": 1210},
"NP02_PreProd": {"08/2018": 812, "09/2018": 1389},
"NP02_Prod": {"08/2018": 3742, "09/2018": 4716},
"FM_Integ": {"08/2018": 1498},
"FM_Prod": {"09/2018": 5686},
}
CUMULATIVE_BUDGET_AARDVARK = {
"02/2018": {"spend": 9857, "cumulative": 9857},
"03/2018": {"spend": 7881, "cumulative": 17738},
"04/2018": {"spend": 14010, "cumulative": 31748},
"05/2018": {"spend": 43510, "cumulative": 75259},
"06/2018": {"spend": 41725, "cumulative": 116_984},
"07/2018": {"spend": 41328, "cumulative": 158_312},
"08/2018": {"spend": 47491, "cumulative": 205_803},
"09/2018": {"spend": 36028, "cumulative": 241_831},
}
CUMULATIVE_BUDGET_BELUGA = {
"08/2018": {"spend": 4838, "cumulative": 4838},
"09/2018": {"spend": 14500, "cumulative": 19338},
}
REPORT_FIXTURE_MAP = {
"Aardvark": {
"cumulative": CUMULATIVE_BUDGET_AARDVARK,
"projects": [
MockProject("LC04", ["Integ", "PreProd", "Prod"]),
MockProject("SF18", ["Integ", "PreProd", "Prod"]),
MockProject("Canton", ["Prod"]),
MockProject("BD04", ["Integ", "PreProd"]),
MockProject("SCV18", ["Dev"]),
MockProject(
"Crown",
[
"CR Portal Dev",
"CR Staging",
"CR Portal Test 1",
"Jewels Prod",
"Jewels Dev",
],
),
],
"budget": 500_000,
},
"Beluga": {
"cumulative": CUMULATIVE_BUDGET_BELUGA,
"projects": [
MockProject("NP02", ["Integ", "PreProd", "NP02_Prod"]),
MockProject("FM", ["Integ", "Prod"]),
],
"budget": 70000,
},
}
def _sum_monthly_spend(self, data):
return sum(
[
spend
for project in data
for env in project.environments
for spend in self.MONTHLY_SPEND_BY_ENVIRONMENT[env.id].values()
]
)
def get_budget(self, workspace):
if workspace.name in self.REPORT_FIXTURE_MAP:
return self.REPORT_FIXTURE_MAP[workspace.name]["budget"]
elif workspace.request and workspace.legacy_task_order:
return workspace.legacy_task_order.budget
return 0
def get_total_spending(self, workspace):
if workspace.name in self.REPORT_FIXTURE_MAP:
return self._sum_monthly_spend(
self.REPORT_FIXTURE_MAP[workspace.name]["projects"]
)
return 0
def _rollup_project_totals(self, data):
project_totals = {}
for project, environments in data.items():
project_spend = [
(month, spend)
for env in environments.values()
if env
for month, spend in env.items()
]
project_totals[project] = {
month: sum([spend[1] for spend in spends])
for month, spends in groupby(sorted(project_spend), lambda x: x[0])
}
return project_totals
def _rollup_workspace_totals(self, project_totals):
monthly_spend = [
(month, spend)
for project in project_totals.values()
for month, spend in project.items()
]
workspace_totals = {}
for month, spends in groupby(sorted(monthly_spend), lambda m: m[0]):
workspace_totals[month] = sum([spend[1] for spend in spends])
return workspace_totals
def monthly_totals_for_environment(self, environment_id):
"""Return the monthly totals for the specified environment.
Data should be in the format of a dictionary with the month as the key
and the spend in that month as the value. For example:
{ "01/2018": 79.85, "02/2018": 86.54 }
"""
return self.MONTHLY_SPEND_BY_ENVIRONMENT.get(environment_id, {})
def monthly_totals(self, workspace):
"""Return month totals rolled up by environment, project, and workspace.
Data should returned with three top level keys, "workspace", "projects",
and "environments".
The "projects" key will have budget data per month for each project,
The "environments" key will have budget data for each environment.
The "workspace" key will be total monthly spending for the workspace.
For example:
{
"environments": { "X-Wing": { "Prod": { "01/2018": 75.42 } } },
"projects": { "X-Wing": { "01/2018": 75.42 } },
"workspace": { "01/2018": 75.42 },
}
"""
projects = workspace.projects
if workspace.name in self.REPORT_FIXTURE_MAP:
projects = self.REPORT_FIXTURE_MAP[workspace.name]["projects"]
environments = {
project.name: {
env.name: self.monthly_totals_for_environment(env.id)
for env in project.environments
}
for project in projects
}
project_totals = self._rollup_project_totals(environments)
workspace_totals = self._rollup_workspace_totals(project_totals)
return {
"environments": environments,
"projects": project_totals,
"workspace": workspace_totals,
}
def cumulative_budget(self, workspace):
if workspace.name in self.REPORT_FIXTURE_MAP:
budget_months = self.REPORT_FIXTURE_MAP[workspace.name]["cumulative"]
else:
budget_months = {}
this_year = pendulum.now().year
all_months = OrderedDict()
for m in range(1, 13):
month_str = "{month:02d}/{year}".format(month=m, year=this_year)
all_months[month_str] = budget_months.get(month_str, None)
return {"months": all_months}

View File

@ -30,3 +30,7 @@ class UnauthenticatedError(Exception):
@property @property
def message(self): def message(self):
return str(self) return str(self)
class UploadError(Exception):
pass

View File

@ -1,249 +1,17 @@
from itertools import groupby from flask import current_app
from collections import OrderedDict
import pendulum
MONTHLY_SPEND_AARDVARK = {
"LC04": {
"Integ": {
"02/2018": 284,
"03/2018": 1210,
"04/2018": 1430,
"05/2018": 1366,
"06/2018": 1169,
"07/2018": 991,
"08/2018": 978,
"09/2018": 737,
},
"PreProd": {
"02/2018": 812,
"03/2018": 1389,
"04/2018": 1425,
"05/2018": 1306,
"06/2018": 1112,
"07/2018": 936,
"08/2018": 921,
"09/2018": 694,
},
"Prod": {
"02/2018": 1742,
"03/2018": 1716,
"04/2018": 1866,
"05/2018": 1809,
"06/2018": 1839,
"07/2018": 1633,
"08/2018": 1654,
"09/2018": 1103,
},
},
"SF18": {
"Integ": {
"04/2018": 1498,
"05/2018": 1400,
"06/2018": 1394,
"07/2018": 1171,
"08/2018": 1200,
"09/2018": 963,
},
"PreProd": {
"04/2018": 1780,
"05/2018": 1667,
"06/2018": 1703,
"07/2018": 1474,
"08/2018": 1441,
"09/2018": 933,
},
"Prod": {
"04/2018": 1686,
"05/2018": 1779,
"06/2018": 1792,
"07/2018": 1570,
"08/2018": 1539,
"09/2018": 986,
},
},
"Canton": {
"Prod": {
"05/2018": 28699,
"06/2018": 26766,
"07/2018": 22619,
"08/2018": 24090,
"09/2018": 16719,
}
},
"BD04": {
"Integ": {},
"PreProd": {
"02/2018": 7019,
"03/2018": 3004,
"04/2018": 2691,
"05/2018": 2901,
"06/2018": 3463,
"07/2018": 3314,
"08/2018": 3432,
"09/2018": 723,
},
},
"SCV18": {"Dev": {"05/2019": 9797}},
"Crown": {
"CR Portal Dev": {
"03/2018": 208,
"04/2018": 457,
"05/2018": 671,
"06/2018": 136,
"07/2018": 1524,
"08/2018": 2077,
"09/2018": 1858,
},
"CR Staging": {
"03/2018": 208,
"04/2018": 457,
"05/2018": 671,
"06/2018": 136,
"07/2018": 1524,
"08/2018": 2077,
"09/2018": 1858,
},
"CR Portal Test 1": {"07/2018": 806, "08/2018": 1966, "09/2018": 2597},
"Jewels Prod": {"07/2018": 806, "08/2018": 1966, "09/2018": 2597},
"Jewels Dev": {
"03/2018": 145,
"04/2018": 719,
"05/2018": 1243,
"06/2018": 2214,
"07/2018": 2959,
"08/2018": 4151,
"09/2018": 4260,
},
},
}
CUMULATIVE_BUDGET_AARDVARK = {
"02/2018": {"spend": 9857, "cumulative": 9857},
"03/2018": {"spend": 7881, "cumulative": 17738},
"04/2018": {"spend": 14010, "cumulative": 31748},
"05/2018": {"spend": 43510, "cumulative": 75259},
"06/2018": {"spend": 41725, "cumulative": 116_984},
"07/2018": {"spend": 41328, "cumulative": 158_312},
"08/2018": {"spend": 47491, "cumulative": 205_803},
"09/2018": {"spend": 36028, "cumulative": 241_831},
}
MONTHLY_SPEND_BELUGA = {
"NP02": {
"Integ": {"08/2018": 284, "09/2018": 1210},
"PreProd": {"08/2018": 812, "09/2018": 1389},
"Prod": {"08/2018": 3742, "09/2018": 4716},
},
"FM": {"Integ": {"08/2018": 1498}, "Prod": {"09/2018": 5686}},
}
CUMULATIVE_BUDGET_BELUGA = {
"08/2018": {"spend": 4838, "cumulative": 4838},
"09/2018": {"spend": 14500, "cumulative": 19338},
}
REPORT_FIXTURE_MAP = {
"Aardvark": {
"cumulative": CUMULATIVE_BUDGET_AARDVARK,
"monthly": MONTHLY_SPEND_AARDVARK,
"budget": 500_000,
},
"Beluga": {
"cumulative": CUMULATIVE_BUDGET_BELUGA,
"monthly": MONTHLY_SPEND_BELUGA,
"budget": 70000,
},
}
def _sum_monthly_spend(data):
return sum(
[
spend
for project in data.values()
for env in project.values()
for spend in env.values()
]
)
def _derive_project_totals(data):
project_totals = {}
for project, environments in data.items():
project_spend = [
(month, spend)
for env in environments.values()
for month, spend in env.items()
]
project_totals[project] = {
month: sum([spend[1] for spend in spends])
for month, spends in groupby(sorted(project_spend), lambda x: x[0])
}
return project_totals
def _derive_workspace_totals(project_totals):
monthly_spend = [
(month, spend)
for project in project_totals.values()
for month, spend in project.items()
]
workspace_totals = {}
for month, spends in groupby(sorted(monthly_spend), lambda m: m[0]):
workspace_totals[month] = sum([spend[1] for spend in spends])
return workspace_totals
class Reports: class Reports:
@classmethod @classmethod
def workspace_totals(cls, workspace): def workspace_totals(cls, workspace):
if workspace.name in REPORT_FIXTURE_MAP: budget = current_app.csp.reports.get_budget(workspace)
budget = REPORT_FIXTURE_MAP[workspace.name]["budget"] spent = current_app.csp.reports.get_total_spending(workspace)
spent = _sum_monthly_spend(REPORT_FIXTURE_MAP[workspace.name]["monthly"])
elif workspace.request and workspace.request.legacy_task_order:
ws_to = workspace.request.legacy_task_order
budget = ws_to.budget
# spent will be derived from CSP data
spent = 0
else:
budget = 0
spent = 0
return {"budget": budget, "spent": spent} return {"budget": budget, "spent": spent}
@classmethod @classmethod
def monthly_totals(cls, workspace): def monthly_totals(cls, workspace):
if workspace.name in REPORT_FIXTURE_MAP: return current_app.csp.reports.monthly_totals(workspace)
environments = REPORT_FIXTURE_MAP[workspace.name]["monthly"]
else:
environments = {
project.name: {env.name: {} for env in project.environments}
for project in workspace.projects
}
project_totals = _derive_project_totals(environments)
workspace_totals = _derive_workspace_totals(project_totals)
return {
"environments": environments,
"projects": project_totals,
"workspace": workspace_totals,
}
@classmethod @classmethod
def cumulative_budget(cls, workspace): def cumulative_budget(cls, workspace):
if workspace.name in REPORT_FIXTURE_MAP: return current_app.csp.reports.cumulative_budget(workspace)
budget_months = REPORT_FIXTURE_MAP[workspace.name]["cumulative"]
else:
budget_months = {}
this_year = pendulum.now().year
all_months = OrderedDict()
for m in range(1, 13):
month_str = "{month:02d}/{year}".format(month=m, year=this_year)
all_months[month_str] = budget_months.get(month_str, None)
return {"months": all_months}

View File

@ -5,8 +5,7 @@ from flask import current_app as app
from atst.models import Base, types, mixins from atst.models import Base, types, mixins
from atst.database import db from atst.database import db
from atst.uploader import UploadError from atst.domain.exceptions import NotFoundError, UploadError
from atst.domain.exceptions import NotFoundError
class AttachmentError(Exception): class AttachmentError(Exception):
@ -25,12 +24,12 @@ class Attachment(Base, mixins.TimestampsMixin):
@classmethod @classmethod
def attach(cls, fyle, resource=None, resource_id=None): def attach(cls, fyle, resource=None, resource_id=None):
try: try:
filename, object_name = app.uploader.upload(fyle) object_name = app.csp.files.upload(fyle)
except UploadError as e: except UploadError as e:
raise AttachmentError("Could not add attachment. " + str(e)) raise AttachmentError("Could not add attachment. " + str(e))
attachment = Attachment( attachment = Attachment(
filename=filename, filename=fyle.filename,
object_name=object_name, object_name=object_name,
resource=resource, resource=resource,
resource_id=resource_id, resource_id=resource_id,

View File

@ -71,7 +71,7 @@ def task_order_pdf_download(request_id):
request = Requests.get(g.current_user, request_id) request = Requests.get(g.current_user, request_id)
if request.legacy_task_order and request.legacy_task_order.pdf: if request.legacy_task_order and request.legacy_task_order.pdf:
pdf = request.legacy_task_order.pdf pdf = request.legacy_task_order.pdf
generator = app.uploader.download_stream(pdf.object_name) generator = app.csp.files.download(pdf.object_name)
return Response( return Response(
generator, generator,
headers={ headers={

View File

@ -3,3 +3,4 @@ ENVIRONMENT = test
PGDATABASE = atat_test PGDATABASE = atat_test
CRL_DIRECTORY = tests/fixtures/crl CRL_DIRECTORY = tests/fixtures/crl
WTF_CSRF_ENABLED = false WTF_CSRF_ENABLED = false
STORAGE_PROVIDER=LOCAL

View File

@ -2,27 +2,23 @@ import os
import pytest import pytest
from werkzeug.datastructures import FileStorage from werkzeug.datastructures import FileStorage
from atst.uploader import Uploader, UploadError from atst.domain.csp.files import RackspaceFileProvider
from atst.domain.exceptions import UploadError
from tests.mocks import PDF_FILENAME from tests.mocks import PDF_FILENAME
@pytest.fixture(scope="function")
def upload_dir(tmpdir):
return tmpdir.mkdir("uploads")
@pytest.fixture @pytest.fixture
def uploader(upload_dir): def uploader(app):
return Uploader("LOCAL", container=upload_dir) return RackspaceFileProvider(app)
NONPDF_FILENAME = "tests/fixtures/disa-pki.html" NONPDF_FILENAME = "tests/fixtures/disa-pki.html"
def test_upload(uploader, upload_dir, pdf_upload): def test_upload(app, uploader, pdf_upload):
filename, object_name = uploader.upload(pdf_upload) object_name = uploader.upload(pdf_upload)
assert filename == PDF_FILENAME upload_dir = app.config["STORAGE_CONTAINER"]
assert os.path.isfile(os.path.join(upload_dir, object_name)) assert os.path.isfile(os.path.join(upload_dir, object_name))
@ -33,17 +29,18 @@ def test_upload_fails_for_non_pdfs(uploader):
uploader.upload(fs) uploader.upload(fs)
def test_download_stream(upload_dir, uploader, pdf_upload): def test_download(app, uploader, pdf_upload):
# write pdf content to upload file storage and make sure it is flushed to # write pdf content to upload file storage and make sure it is flushed to
# disk # disk
pdf_upload.seek(0) pdf_upload.seek(0)
pdf_content = pdf_upload.read() pdf_content = pdf_upload.read()
pdf_upload.close() pdf_upload.close()
upload_dir = app.config["STORAGE_CONTAINER"]
full_path = os.path.join(upload_dir, "abc") full_path = os.path.join(upload_dir, "abc")
with open(full_path, "wb") as output_file: with open(full_path, "wb") as output_file:
output_file.write(pdf_content) output_file.write(pdf_content)
output_file.flush() output_file.flush()
stream = uploader.download_stream("abc") stream = uploader.download("abc")
stream_content = b"".join([b for b in stream]) stream_content = b"".join([b for b in stream])
assert pdf_content == stream_content assert pdf_content == stream_content

View File

@ -9,6 +9,7 @@ from tests.mocks import PDF_FILENAME
def test_attach(pdf_upload): def test_attach(pdf_upload):
attachment = Attachment.attach(pdf_upload) attachment = Attachment.attach(pdf_upload)
assert attachment.filename == PDF_FILENAME assert attachment.filename == PDF_FILENAME
assert attachment.object_name is not None
def test_attach_raises(): def test_attach_raises():