Add tests for S3Cache

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()