Fix api stream (#743)

* Fix api stream

* Fix lint

discovery-provider/src/api/v1/tracks.py

@@ -47,6 +47,9 @@ def get(self, track_id):
         return success_response(single_track)
 
 
+def tranform_stream_cache(stream_url):
+    return redirect(stream_url)
+
 @ns.route("/<string:track_id>/stream")
 class TrackStream(Resource):
     @record_metrics
@@ -61,7 +64,7 @@ class TrackStream(Resource):
             500: 'Server error'
         }
     )
-    @cache(ttl_sec=5)
+    @cache(ttl_sec=5, transform=tranform_stream_cache)
     def get(self, track_id):
         """
         Get the track's streamable mp3 file.
@@ -80,8 +83,8 @@ def get(self, track_id):
 
         primary_node = creator_nodes[0]
         stream_url = urljoin(primary_node, 'tracks/stream/{}'.format(track_id))
-        return redirect(stream_url)
 
+        return stream_url
 
 track_search_result = make_response(
     "track_search", ns, fields.List(fields.Nested(track)))

discovery-provider/src/utils/redis_cache.py

@@ -1,33 +1,45 @@
-import logging # pylint: disable=C0302
+import logging  # pylint: disable=C0302
 import functools
 import json
-import redis
 from flask.json import dumps
 from flask.globals import request
-from src.utils.config import shared_config
+from src.utils import redis_connection
 from src.utils.query_params import stringify_query_params
 logger = logging.getLogger(__name__)
 
-REDIS_URL = shared_config["redis"]["url"]
-REDIS = redis.Redis.from_url(url=REDIS_URL)
-
 # Redis Key Convention:
 # API_V1:path:queryparams
 
 cache_prefix = "API_V1_ROUTE"
 default_ttl_sec = 60
 
+
 def extract_key():
     path = request.path
     req_args = request.args.items()
     req_args = stringify_query_params(req_args)
     key = f"{cache_prefix}:{path}:{req_args}"
     return key
 
+
 def cache(**kwargs):
     """
     Cache decorator.
-    Should be called with `@cache(ttl_sec=123)`
+    Should be called with `@cache(ttl_sec=123, transform=transform_response)`
+
+    Arguments:
+        ttl_sec: optional,number The time in seconds to cache the response if
+            status code < 400
+        transform: optional,func The transform function of the wrapped function
+            to convert the function response to request response
+
+    Usage Notes:
+        If the wrapper function returns a tuple, the transform function will not
+        be run on the response. The first item of the tuple must be serializable.
+
+        If the wrapper function returns a single response, the transform function
+        must be passed to the decorator. The wrapper function response must be
+        serializable.
 
     Decorators in Python are just higher-order-functions that accept a function
     as a single parameter, and return a function that wraps the input function.
@@ -41,20 +53,31 @@ def cache(**kwargs):
     `func` rather than `inner_wrap`.
     """
     ttl_sec = kwargs["ttl_sec"] if "ttl_sec" in kwargs else default_ttl_sec
+    transform = kwargs["transform"] if "transform" in kwargs else None
+    redis = redis_connection.get_redis()
+
     def outer_wrap(func):
         @functools.wraps(func)
         def inner_wrap(*args, **kwargs):
             key = extract_key()
-            cached_resp = REDIS.get(key)
+            cached_resp = redis.get(key)
 
             if cached_resp:
                 deserialized = json.loads(cached_resp)
+                if transform is not None:
+                    return transform(deserialized)
                 return deserialized, 200
 
-            resp, status = func(*args, **kwargs)
-            if status == 200:
-                serialized = dumps(resp)
-                REDIS.set(key, serialized, ttl_sec)
-            return resp, status
+            response = func(*args, **kwargs)
+
+            if len(response) == 2:
+                resp, status_code = response
+                if status_code < 400:
+                    serialized = dumps(resp)
+                    redis.set(key, serialized, ttl_sec)
+                return resp, status_code
+            serialized = dumps(response)
+            redis.set(key, serialized, ttl_sec)
+            return transform(response)
         return inner_wrap
     return outer_wrap

discovery-provider/src/utils/redis_cache_test.py

@@ -0,0 +1,65 @@
+import json
+from time import sleep
+from unittest.mock import patch
+from src.utils.redis_cache import cache
+
+
+def test_cache(redis_mock):
+    """Test that the redis cache decorator works"""
+
+    @patch('src.utils.redis_cache.extract_key')
+    def get_mock_cache(extract_key):
+
+        mock_key_1 = 'mock_key'
+        extract_key.return_value = mock_key_1
+
+        # Test a mock function returning two items
+        @cache(ttl_sec=1)
+        def mock_func():
+            return {'name': 'joe'}, 200
+
+        res = mock_func()
+        assert res[0] == {'name': 'joe'}
+        assert res[1] == 200
+
+        cached_resp = redis_mock.get(mock_key_1)
+        deserialized = json.loads(cached_resp)
+        assert deserialized == {'name': 'joe'}
+
+        # This should call the function and return the cached response
+        res = mock_func()
+        assert res[0] == {'name': 'joe'}
+        assert res[1] == 200
+
+        # Sleep to wait for the cache to expire
+        sleep(1)
+
+        cached_resp = redis_mock.get(mock_key_1)
+        assert cached_resp is None
+
+        # Test the single response
+        def transform(input):
+            return {'music': input}
+
+        @cache(ttl_sec=1, transform=transform)
+        def mock_func_transform():
+            return 'audius'
+
+        res = mock_func_transform()
+        assert res == {'music': 'audius'}
+
+        cached_resp = redis_mock.get(mock_key_1)
+        deserialized = json.loads(cached_resp)
+        assert deserialized == 'audius'
+
+        # This should call the function and return the cached response
+        res = mock_func_transform()
+        assert res == {'music': 'audius'}
+
+        # Sleep to wait for the cache to expire
+        sleep(1)
+
+        cached_resp = redis_mock.get(mock_key_1)
+        assert cached_resp is None
+
+    get_mock_cache() # pylint: disable=no-value-for-parameter