From a7770d8a6a1b5d8cfe4e1c688b09f529d3cd169c Mon Sep 17 00:00:00 2001
From: graham-dds <graham.beckley-ctr@friends.dds.mil>
Date: Thu, 30 Jan 2020 13:27:34 -0500
Subject: [PATCH] Add method to download TOs from Azure

---
 atst/domain/csp/file_uploads.py         | 36 +++++++++++++++++++++----
 tests/domain/cloud/test_file_uploads.py | 22 +++++++++++++++
 2 files changed, 53 insertions(+), 5 deletions(-)
 create mode 100644 tests/domain/cloud/test_file_uploads.py

diff --git a/atst/domain/csp/file_uploads.py b/atst/domain/csp/file_uploads.py
index 7916797d..ce18df63 100644
--- a/atst/domain/csp/file_uploads.py
+++ b/atst/domain/csp/file_uploads.py
@@ -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,
+        }
diff --git a/tests/domain/cloud/test_file_uploads.py b/tests/domain/cloud/test_file_uploads.py
new file mode 100644
index 00000000..9b1b0932
--- /dev/null
+++ b/tests/domain/cloud/test_file_uploads.py
@@ -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"