Improved testing!

Author: jhodges10

frameioclient/client.py

@@ -351,8 +351,7 @@ def copy_asset(self, destination_folder_id, **kwargs):
     return self._api_call('post', endpoint, kwargs)
 
   def bulk_copy_assets(self, destination_folder_id, asset_list=[], copy_comments=False):
-    """
-    Bulk copy assets
+    """Bulk copy assets
 
     :Args:
       destination_folder_id (string): The id of the folder you want to copy into.

frameioclient/download.py

@@ -12,21 +12,24 @@ def __init__(self, asset, download_folder):
   def get_download_key(self):
     try:
       url = self.asset['original']
-    except KeyError:
+    except KeyError as e:
       if self.asset['is_session_watermarked'] == True:
         resolution_list = list()
-        for resolution_key, download_url in sorted(self.asset['downloads'].items()):
-          resolution = resolution_key.split("_")[1] # Grab the item at index 1 (resolution)
-          try:
-            resolution = int(resolution)
-          except ValueError:
-            continue
-
-          if download_url is not None:
-            resolution_list.append(download_url)
-
-        # Grab the highest resolution (first item) now
-        url = resolution_list[0]
+        try:
+          for resolution_key, download_url in sorted(self.asset['downloads'].items()):
+            resolution = resolution_key.split("_")[1] # Grab the item at index 1 (resolution)
+            try:
+              resolution = int(resolution)
+            except ValueError:
+              continue
+
+            if download_url is not None:
+              resolution_list.append(download_url)
+
+          # Grab the highest resolution (first item) now
+          url = resolution_list[0]
+        except KeyError:
+          raise DownloadException
       else:
         raise DownloadException
 

tests/test_frameiodownloader.py

@@ -0,0 +1,56 @@
+import pytest
+from frameioclient.download import FrameioDownloader
+from frameioclient.exceptions import DownloadException
+
+
+regular_asset = {
+    "is_hls_required": False,
+    "is_session_watermarked": False,
+    "downloads": {
+        "h264_720": "some-720-url",
+        "h264_1080_best": "some-1080-url"
+    },
+    "h264_720": "some-720-url",
+    "h264_1080_best": "some-1080-url",
+    "original": "some-original-url",
+    "hls_manifest": "some-hls-url"
+}
+
+watermarked_asset_download_allowed = {
+    "is_hls_required": True,
+    "is_session_watermarked": True,
+    "downloads": {
+        "h264_720": "download-stream-service-url",
+        "h264_1080_best": "download-stream-service-url"
+    },
+    "hls_manifest": "hls-url"
+}
+
+watermarked_asset_no_download = {
+    "is_hls_required": True,
+    "is_session_watermarked": True,
+    "hls_manifest": "hls-url"
+}
+
+no_download_allowed = {
+    "is_hls_required": True,
+    "is_session_watermarked": False,
+    "hls_manifest": "hls-url"
+}
+
+def test_get_download_key_returns_original():
+    url = FrameioDownloader(regular_asset, './').get_download_key()
+    assert url == regular_asset['original']
+
+def test_get_download_key_returns_watermarked_download():
+    url = FrameioDownloader(watermarked_asset_download_allowed, './').get_download_key()
+    assert url == watermarked_asset_download_allowed['downloads']['h264_1080_best']
+
+def test_get_download_key_fails_gracefully_on_watermarked_asset():
+    with pytest.raises(DownloadException):
+        FrameioDownloader(watermarked_asset_no_download, './').get_download_key()
+
+def test_get_download_key_fails_gracefully_when_downloads_disallowed():
+    with pytest.raises(DownloadException):
+        FrameioDownloader(no_download_allowed, './').get_download_key()
+