Added simple mocking test for GCS

This commit is contained in:
Sybren A. Stüvel 2017-03-22 16:43:17 +01:00
parent c02c6d89b0
commit b36dc63335

View File

@ -1,9 +1,10 @@
import typing
from unittest import mock
from pillar.tests import AbstractPillarTest
class LocalStorageBackendTest(AbstractPillarTest):
class AbstractStorageBackendTest(AbstractPillarTest):
def create_test_file(self) -> (typing.IO, bytes):
import io
import secrets
@ -13,6 +14,20 @@ class LocalStorageBackendTest(AbstractPillarTest):
return test_file, file_contents
def assert_valid_file(self, expected_file_contents: bytes, url: str):
resp = self.get(url)
self.assertEqual(200, resp.status_code)
self.assertEqual(len(expected_file_contents), int(resp.headers['Content-Length']))
self.assertEqual(expected_file_contents, resp.data)
class LocalStorageBackendTest(AbstractStorageBackendTest):
def storage_backend(self):
from pillar.api.file_storage_backends import Bucket
return Bucket.for_backend('local')
def test_upload_download(self):
test_file, file_contents = self.create_test_file()
@ -24,19 +39,33 @@ class LocalStorageBackendTest(AbstractPillarTest):
# We should be able to upload the file, and then download it again
# from the URL given by its blob.
blob.create_from_file(test_file, file_size=512, content_type='application/octet-stream')
url = blob.get_url(is_public=True)
resp = self.get(url)
self.assert_valid_file(file_contents, url)
self.assertEqual(200, resp.status_code)
self.assertEqual('512', resp.headers['Content-Length'])
self.assertEqual(file_contents, resp.data)
def test_upload_from_path(self):
import tempfile
import pathlib
def storage_backend(self):
from pillar.api.file_storage_backends import Bucket
test_file, file_contents = self.create_test_file()
with tempfile.TemporaryDirectory() as tmpdir:
test_file_path = pathlib.Path(tmpdir) / 'testfile.bin'
return Bucket.for_backend('local')
with test_file_path.open('wb') as outfile:
outfile.write(file_contents)
with self.app.test_request_context():
bucket_class = self.storage_backend()
bucket = bucket_class('buckettest')
blob = bucket.blob('somefile.bin')
# We should be able to upload the file, and then download it again
# from the URL given by its blob, even after the original file was removed.
blob.upload_from_path(test_file_path, content_type='application/octet-stream')
url = blob.get_url(is_public=True)
self.assertFalse(test_file_path.exists())
self.assert_valid_file(file_contents, url)
def test_copy_to_bucket(self):
from bson import ObjectId
@ -64,3 +93,45 @@ class LocalStorageBackendTest(AbstractPillarTest):
self.assertEqual(200, resp.status_code)
self.assertEqual('512', resp.headers['Content-Length'])
self.assertEqual(file_contents, resp.data)
class MockedGoogleCloudStorageTest(AbstractStorageBackendTest):
def storage_backend(self):
from pillar.api.file_storage_backends import Bucket
return Bucket.for_backend('gcs')
def test_file_upload(self):
import pillar.api.file_storage_backends.gcs as gcs
from gcloud.storage import Client, Bucket, Blob
# Set up mock GCS client
mock_gcs_client = gcs.gcs = mock.MagicMock(name='mock_gcs_client', autospec=Client)
mock_bucket = mock.MagicMock(name='mock_bucket', autospec=Bucket)
mock_blob = mock.MagicMock(name='mock_blob', autospec=Blob)
mock_gcs_client.get_bucket.return_value = mock_bucket
mock_bucket.blob.return_value = mock_blob
mock_blob.public_url = '/path/to/somefile.bin'
mock_blob.size = 318
test_file, file_contents = self.create_test_file()
with self.app.test_request_context():
bucket_class = self.storage_backend()
bucket = bucket_class('buckettest')
blob = bucket.blob('somefile.bin')
# We should be able to upload the file, and then download it again
# from the URL given by its blob.
blob.create_from_file(test_file, file_size=512, content_type='application/octet-stream')
url = blob.get_url(is_public=True)
self.assertIn('somefile.bin', url)
# Google-reported size should take precedence over reality.
self.assertEqual(318, blob.size)
mock_blob.upload_from_file.assert_called_with(test_file,
size=512,
content_type='application/octet-stream')
mock_blob.reload.assert_called_once()