Add method to download TOs from Azure

This commit is contained in:
graham-dds 2020-01-30 13:27:34 -05:00 committed by leigh-mil
parent 0e49d038be
commit a7770d8a6a
2 changed files with 53 additions and 5 deletions

View File

@ -4,14 +4,17 @@ from uuid import uuid4
class Uploader:
def generate_token(self):
pass
raise NotImplementedError()
def generate_download_link(self, object_name, filename) -> (dict, str):
pass
raise NotImplementedError()
def object_name(self) -> str:
return str(uuid4())
def download_task_order(self, object_name):
raise NotImplementedError()
class MockUploader(Uploader):
def __init__(self, config):
@ -23,6 +26,13 @@ class MockUploader(Uploader):
def generate_download_link(self, object_name, filename):
return ""
def download_task_order(self, object_name):
with open("tests/fixtures/sample.pdf", "rb") as some_bytes:
return {
"name": object_name,
"content": some_bytes,
}
class AzureUploader(Uploader):
def __init__(self, config):
@ -32,10 +42,12 @@ class AzureUploader(Uploader):
self.timeout = timedelta(seconds=config["PERMANENT_SESSION_LIFETIME"])
from azure.storage.common import CloudStorageAccount
from azure.storage.blob import BlobPermissions
from azure.storage.blob import BlobSasPermissions
from azure.storage.blob.blockblobservice import BlockBlobService
self.CloudStorageAccount = CloudStorageAccount
self.BlobPermissions = BlobPermissions
self.BlobSasPermissions = BlobSasPermissions
self.BlockBlobService = BlockBlobService
def get_token(self):
"""
@ -53,7 +65,7 @@ class AzureUploader(Uploader):
sas_token = bbs.generate_blob_shared_access_signature(
self.container_name,
object_name,
permission=self.BlobPermissions.CREATE,
permission=self.BlobSasPermissions(create=True),
expiry=datetime.utcnow() + self.timeout,
protocol="https",
)
@ -75,3 +87,17 @@ class AzureUploader(Uploader):
return bbs.make_blob_url(
self.container_name, object_name, protocol="https", sas_token=sas_token
)
def download_task_order(self, object_name):
block_blob_service = self.BlockBlobService(
account_name=self.account_name, account_key=self.storage_key
)
# TODO: We should downloading errors more gracefully
# - what happens when we try to request a TO that doesn't exist?
b = block_blob_service.get_blob_to_bytes(
container_name=self.container_name, blob_name=object_name,
)
return {
"name": b.name,
"content": b.content,
}

View File

@ -0,0 +1,22 @@
from atst.domain.csp.file_uploads import AzureUploader
from azure.storage.blob.models import Blob
class MockBlockBlobService(object):
def __init__(self, exception=None, **kwargs):
self.exception = exception
def get_blob_to_bytes(self, blob_name="test.pdf", **kwargs):
if self.exception:
raise self.exception
else:
return Blob(name=blob_name, content=b"mock content")
def test_download_task_order_success(app, monkeypatch):
uploader = AzureUploader(config=app.config)
uploader.BlockBlobService = MockBlockBlobService
task_order = uploader.download_task_order("test.pdf")
assert task_order["name"] == "test.pdf"
assert task_order["content"] == b"mock content"