Add tests for S3Cache

superset/assets/version_info.json

@@ -1 +1 @@
-{"GIT_SHA": "0f7189b859f4a782fd43af694012029645f81b44", "version": "0.15.4"}
+{"GIT_SHA": "1e94498d9d548cbea6466a45dafa3b919c65bd1f", "version": "0.15.4"}

superset/results_backends.py

@@ -35,7 +35,7 @@ def __init__(self, default_timeout=300):
         self.default_timeout = default_timeout
 
         self.s3_client = boto3.client('s3')
-        self.bucket = self.s3_resource.Bucket(config.get('S3_CACHE_BUCKET'))
+        self.bucket = config.get('S3_CACHE_BUCKET')
         self.key_prefix = config.get('S3_CACHE_KEY_PREFIX')
 
     def get(self, key):
@@ -140,7 +140,7 @@ def _full_s3_key(self, key):
     def _key_exists(self, key):
         """Determine whether the given key exists in the bucket."""
         try:
-            response = self.s3_client.head_object(
+            self.s3_client.head_object(
                 Bucket=self.bucket,
                 Key=self._full_s3_key(key)
             )

tests/results_backends_tests.py

@@ -0,0 +1,120 @@
+import cPickle
+import mock
+
+from superset import app, results_backends
+from .base_tests import SupersetTestCase
+
+app.config['S3_CACHE_BUCKET'] = 'test-bucket'
+app.config['S3_CACHE_KEY_PREFIX'] = 'test-prefix/'
+
+
+class ResultsBackendsTests(SupersetTestCase):
+    requires_examples = False
+
+    @mock.patch('boto3.client')
+    def setUp(self, mock_boto3_client):
+        self.mock_boto3_client = mock_boto3_client
+        self.mock_s3_client = mock.MagicMock()
+
+        self.mock_boto3_client.return_value = self.mock_s3_client
+
+        self.s3_cache = results_backends.S3Cache()
+        self.s3_cache._key_exists = ResultsBackendsTests._mock_key_exists
+
+    @staticmethod
+    def _mock_download_fileobj(bucket, key, value_file):
+        value_file.write(cPickle.dumps('%s:%s' % (bucket, key)))
+
+    @staticmethod
+    def _mock_key_exists(key):
+        return key == 'test-key'
+
+    def test_s3_cache_initilization(self):
+        self.mock_boto3_client.assert_called_with('s3')
+
+    def test_s3_cache_set(self):
+        result = self.s3_cache.set('test-key', 'test-value')
+
+        self.assertTrue(result)
+        self.mock_s3_client.upload_fileobj.assert_called_once()
+
+        call_args = self.mock_s3_client.upload_fileobj.call_args_list[0][0]
+
+        self.assertEquals(cPickle.loads(call_args[0].getvalue()), 'test-value')
+        self.assertEquals(call_args[1], 'test-bucket')
+        self.assertEquals(call_args[2], 'test-prefix/test-key')
+
+    def test_s3_cache_exception(self):
+        self.mock_s3_client.upload_fileobj.side_effect = Exception('Something bad happened!')
+        result = self.s3_cache.set('test-key', 'test-value')
+
+        self.assertFalse(result)
+        self.mock_s3_client.upload_fileobj.assert_called_once()
+
+    def test_s3_cache_get_exists(self):
+        self.mock_s3_client.download_fileobj.side_effect = (
+            ResultsBackendsTests._mock_download_fileobj)
+        result = self.s3_cache.get('test-key')
+
+        self.assertEquals(result, 'test-bucket:test-prefix/test-key')
+        self.mock_s3_client.download_fileobj.assert_called_once()
+
+    def test_s3_cache_get_does_not_exist(self):
+        result = self.s3_cache.get('test-key2')
+
+        self.assertEquals(result, None)
+        self.assertFalse(self.mock_s3_client.download_fileobj.called)
+
+    def test_s3_cache_get_s3_exception(self):
+        self.mock_s3_client.download_fileobj.side_effect = Exception('Something bad happened')
+        result = self.s3_cache.get('test-key')
+
+        self.assertEquals(result, None)
+        self.mock_s3_client.download_fileobj.assert_called_once()
+
+    def test_s3_cache_delete_exists(self):
+        result = self.s3_cache.delete('test-key')
+
+        self.assertTrue(result)
+        self.mock_s3_client.delete_objects.assert_called_once_with(
+            Bucket='test-bucket',
+            Delete={'Objects': [{'Key': 'test-prefix/test-key'}]}
+        )
+
+    def test_s3_cache_delete_does_not_exist(self):
+        result = self.s3_cache.delete('test-key2')
+
+        self.assertFalse(result)
+        self.assertFalse(self.mock_s3_client.delete_objects.called)
+
+    def test_s3_cache_delete_exception(self):
+        self.mock_s3_client.delete_objects.side_effect = Exception('Something bad happened')
+        result = self.s3_cache.delete('test-key')
+
+        self.assertFalse(result)
+        self.mock_s3_client.delete_objects.assert_called_once()
+
+    def test_s3_cache_add_exists(self):
+        result = self.s3_cache.add('test-key', 'test-value')
+
+        self.assertFalse(result)
+        self.assertFalse(self.mock_s3_client.upload_fileobj.called)
+
+    def test_s3_cache_add_does_not_exist(self):
+        result = self.s3_cache.add('test-key2', 'test-value')
+
+        self.assertTrue(result)
+        self.mock_s3_client.upload_fileobj.assert_called_once()
+
+        call_args = self.mock_s3_client.upload_fileobj.call_args_list[0][0]
+
+        self.assertEquals(cPickle.loads(call_args[0].getvalue()), 'test-value')
+        self.assertEquals(call_args[1], 'test-bucket')
+        self.assertEquals(call_args[2], 'test-prefix/test-key2')
+
+    def test_s3_cache_add_exception(self):
+        self.mock_s3_client.upload_fileobj.side_effect = Exception('Something bad happened')
+        result = self.s3_cache.add('test-key2', 'test-value')
+
+        self.assertFalse(result)
+        self.mock_s3_client.upload_fileobj.assert_called_once()