From 608b12bd8ad3efa513b6bd87753e104b11a1d81e Mon Sep 17 00:00:00 2001
From: Trent Smith <trent.smith@chanzuckerberg.com>
Date: Tue, 30 Nov 2021 09:42:58 -0800
Subject: [PATCH] Fix: return response caching

- add check when s3_uri is None
- add typing
- add additional testing
---
 server/app/api/util.py                       |  3 +-
 server/data_common/dataset_metadata.py       |  6 ++-
 server/tests/unit/common/apis/test_api_v3.py | 53 ++++++++------------
 3 files changed, 26 insertions(+), 36 deletions(-)

diff --git a/server/app/api/util.py b/server/app/api/util.py
index 9269d9f71..0c6bd943b 100644
--- a/server/app/api/util.py
+++ b/server/app/api/util.py
@@ -6,7 +6,8 @@
 
 def get_dataset_artifact_s3_uri(url_dataroot: str = None, dataset_id: str = None):
     dataset_artifact_s3_uri = get_dataset_metadata(url_dataroot, dataset_id, current_app.app_config)["s3_uri"]
-    dataset_artifact_s3_uri = dataset_artifact_s3_uri.rstrip("/")  # remove trailing slash for cloudfront caching.
+    if dataset_artifact_s3_uri:
+        dataset_artifact_s3_uri = dataset_artifact_s3_uri.rstrip("/")  # remove trailing slash for cloudfront caching.
     return dataset_artifact_s3_uri
 
 
diff --git a/server/data_common/dataset_metadata.py b/server/data_common/dataset_metadata.py
index 6bc9dbecf..06e02f496 100644
--- a/server/data_common/dataset_metadata.py
+++ b/server/data_common/dataset_metadata.py
@@ -1,5 +1,7 @@
 import json
 import logging
+from typing import Union
+
 import requests
 from flask import current_app
 
@@ -27,7 +29,7 @@ def request_dataset_metadata_from_data_portal(data_portal_api_base: str, explore
         return None
 
 
-def infer_dataset_s3_uri(server_config: ServerConfig, dataset_root: str, dataset_id: str):
+def infer_dataset_s3_uri(server_config: ServerConfig, dataset_root: str, dataset_id: str) -> Union[str, None]:
     """
     Use the dataset_root, dataset_id, and the server config to infer the physical S3 URI for an Explorer dataset
     artifact
@@ -56,7 +58,7 @@ def infer_dataset_s3_uri(server_config: ServerConfig, dataset_root: str, dataset
     return s3_uri
 
 
-def get_dataset_metadata(dataset_root: str, dataset_id: str, app_config: AppConfig):
+def get_dataset_metadata(dataset_root: str, dataset_id: str, app_config: AppConfig) -> dict:
     """
     Given the dataset root and dataset_id and the explorer web base url (from the app_config), returns the metadata for
     the dataset, including the s3 URI of the dataset artifact.
diff --git a/server/tests/unit/common/apis/test_api_v3.py b/server/tests/unit/common/apis/test_api_v3.py
index 1f210b88d..4ee8d326f 100644
--- a/server/tests/unit/common/apis/test_api_v3.py
+++ b/server/tests/unit/common/apis/test_api_v3.py
@@ -731,39 +731,26 @@ def setUpClass(cls):
 
     @patch("server.data_common.dataset_metadata.requests.get")
     def test_get_S3_URI_in_data_portal(self, mock_get):
-        s3_uri = f"{FIXTURES_ROOT}/pbmc3k.cxg"
-        response_body = json.dumps(
-            {
-                "collection_id": "4f098ff4-4a12-446b-a841-91ba3d8e3fa6",
-                "collection_visibility": "PUBLIC",
-                "dataset_id": "2fa37b10-ab4d-49c9-97a8-b4b3d80bf939",
-                "s3_uri": s3_uri,
-                "tombstoned": False,
-            }
-        )
-        mock_get.return_value = MockResponse(body=response_body, status_code=200)
-
-        result = self.client.get(self.url)
-        self.assertEqual(result.status_code, HTTPStatus.OK)
-        self.assertEqual(json.loads(result.data), s3_uri)
-
-    @patch("server.data_common.dataset_metadata.requests.get")
-    def test_get_S3_URI_in_data_portal_format(self, mock_get):
-        s3_uri = f"{FIXTURES_ROOT}/pbmc3k.cxg/"
-        response_body = json.dumps(
-            {
-                "collection_id": "4f098ff4-4a12-446b-a841-91ba3d8e3fa6",
-                "collection_visibility": "PUBLIC",
-                "dataset_id": "2fa37b10-ab4d-49c9-97a8-b4b3d80bf939",
-                "s3_uri": s3_uri,
-                "tombstoned": False,
-            }
-        )
-        mock_get.return_value = MockResponse(body=response_body, status_code=200)
-
-        result = self.client.get(self.url)
-        self.assertEqual(result.status_code, HTTPStatus.OK)
-        self.assertEqual(json.loads(result.data), s3_uri[:-1])
+        test_s3_uris = [
+            (f"{FIXTURES_ROOT}/pbmc3k.cxg", f"{FIXTURES_ROOT}/pbmc3k.cxg"),
+            (f"{FIXTURES_ROOT}/pbmc3k.cxg/", f"{FIXTURES_ROOT}/pbmc3k.cxg"),
+            (None, None),
+        ]
+        test_response_body = {
+            "collection_id": "4f098ff4-4a12-446b-a841-91ba3d8e3fa6",
+            "collection_visibility": "PUBLIC",
+            "dataset_id": "2fa37b10-ab4d-49c9-97a8-b4b3d80bf939",
+            "tombstoned": False,
+        }
+        for actual, expected in test_s3_uris:
+            with self.subTest(actual):
+                test_response_body["s3_uri"] = actual
+                response_body = json.dumps(test_response_body)
+                mock_get.return_value = MockResponse(body=response_body, status_code=200)
+
+                result = self.client.get(self.url)
+                self.assertEqual(result.status_code, HTTPStatus.OK)
+                self.assertEqual(json.loads(result.data), expected)
 
     @patch("server.data_common.dataset_metadata.requests.get")
     def test_get_S3_URI_not_in_data_portal(self, mock_get):