Add method to download TOs from Azure
This commit is contained in:
parent
0e49d038be
commit
a7770d8a6a
@ -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,
|
||||
}
|
||||
|
22
tests/domain/cloud/test_file_uploads.py
Normal file
22
tests/domain/cloud/test_file_uploads.py
Normal 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"
|
Loading…
x
Reference in New Issue
Block a user