def test_list_with_subdir(self, client, _):
blob_root_path = 'project_path/experiment_id/'
gcs_url = 'gs://bucket/' + blob_root_path
dirname = 'model'
obj_mock = mock.Mock()
file_path = dirname + '/' + 'tf.pb'
obj_mock.configure_mock(name=blob_root_path + file_path, size=1)
subdir_mock = mock.Mock()
subdirname = dirname + '/' + 'files'
subdir_mock.configure_mock(prefixes=(blob_root_path + subdirname +
'/', ))
mock_results = mock.MagicMock()
mock_results.configure_mock(pages=[subdir_mock])
mock_results.__iter__.return_value = [obj_mock]
client.return_value.get_bucket.return_value.list_blobs.return_value = mock_results
store = GCSStore()
results = store.list(gcs_url, path=dirname)
blobs = results['blobs']
prefixes = results['prefixes']
assert len(blobs) == 1
assert len(prefixes) == 1
assert blobs[0][0] == file_path
assert blobs[0][1] == obj_mock.size
assert prefixes[0] == subdirname
def test_store_client(self, client, gc_credentials):
client.return_value = 'client'
store = GCSStore(client='foo')
assert store.client == 'foo'
store = GCSStore()
assert store.client == 'client'
assert gc_credentials.call_count == 1
assert client.call_count == 1
def test_existing_object(self, client, _):
test_bucket = 'test_bucket'
test_object = 'test_object'
client.return_value.objects.return_value.get.return_value.execute.return_value = {
"kind": "storage#object",
# the bucket name, object name, and generation number.
"id": "{}/{}/1521132662504504".format(test_bucket, test_object),
"name": test_object,
"bucket": test_bucket,
"generation": "1521132662504504",
"contentType": "text/csv",
"timeCreated": "2018-03-15T16:51:02.502Z",
"updated": "2018-03-15T16:51:02.502Z",
"storageClass": "MULTI_REGIONAL",
"timeStorageClassUpdated": "2018-03-15T16:51:02.502Z",
"size": "89",
"md5Hash": "leYUJBUWrRtks1UeUFONJQ==",
"metadata": {
"md5-hash": "95e614241516ad1b64b3551e50538d25"
},
"crc32c": "xgdNfQ==",
"etag": "CLf4hODk7tkCEAE="
}
response = GCSStore().check_blob(blob=test_object,
bucket_name=test_bucket)
assert response is True
def test_get_bucket(self, client, _):
test_bucket = 'test_bucket'
client.return_value.get_bucket.return_value = {}
response = GCSStore().get_bucket(bucket_name=test_bucket)
assert response == {}
def test_non_existing_object(self, client, _):
test_bucket = 'test_bucket'
test_object = 'test_object'
client.return_value.get_bucket.return_value.get_blob.side_effect = Exception
response = GCSStore().check_blob(blob=test_object, bucket_name=test_bucket)
assert response is False
def test_delete(self, client, _):
test_bucket = 'test_bucket'
test_object = 'test_object'
# Correct file
bucket = mock.MagicMock()
client.return_value.get_bucket.return_value = bucket
bucket.delete_blob.return_value = True
GCSStore().delete(key=test_object, bucket_name=test_bucket)
# Wrong file
bucket = mock.MagicMock()
client.return_value.get_bucket.return_value = bucket
bucket.delete_blob.return_value = True
GCSStore().delete(key=test_object, bucket_name=test_bucket)
def test_parse_gcs_url(self):
# Correct url
gcs_url = 'gs://bucket/path/to/blob'
parsed_url = GCSStore.parse_gcs_url(gcs_url)
assert parsed_url == ('bucket', 'path/to/blob')
# Wrong url
gcs_url = 'gs:/bucket/path/to/blob'
with self.assertRaises(PolyaxonStoresException):
GCSStore.parse_gcs_url(gcs_url)
# Trailing slash
gcs_url = 'gs://bucket/path/to/blob/'
assert GCSStore.parse_gcs_url(gcs_url) == ('bucket', 'path/to/blob/')
# Bucket only
gcs_url = 'gs://bucket/'
assert GCSStore.parse_gcs_url(gcs_url) == ('bucket', '')
def test_get_blob(self, client, _):
test_bucket = 'test_bucket'
test_object = 'test_object'
bucket = mock.MagicMock()
client.return_value.get_bucket.return_value = bucket
bucket.get_blob.return_value = {}
response = GCSStore().get_blob(blob=test_object, bucket_name=test_bucket)
assert response == {}
def test_download(self, client, _):
dirname = tempfile.mkdtemp()
fpath = dirname + '/test.txt'
def mkfile(fname):
return open(fname, 'w')
(client.return_value.get_bucket.return_value.get_blob.return_value.
download_to_filename.side_effect) = mkfile
store = GCSStore()
# Test without basename
gcs_url = 'gs://bucket/path/to/blob.txt'
store.download_file(gcs_url, fpath, use_basename=False)
client.return_value.get_bucket.assert_called_with('bucket')
client.return_value.get_bucket().get_blob.assert_called_with(
'path/to/blob.txt')
client.return_value.get_bucket().get_blob(
).download_to_filename.assert_called_with(fpath)
# Test with basename
gcs_url = 'gs://bucket/path/to/blob.txt'
store.download_file(gcs_url, dirname, use_basename=True)
client.return_value.get_bucket.assert_called_with('bucket')
client.return_value.get_bucket().get_blob.assert_called_with(
'path/to/blob.txt')
client.return_value.get_bucket().get_blob(
).download_to_filename.assert_called_with(dirname + '/blob.txt')
def test_upload(self, client, _):
dirname = tempfile.mkdtemp()
fpath = dirname + '/test.txt'
open(fpath, 'w')
(client.return_value.get_bucket.return_value.blob.return_value.
upload_from_filename.side_effect) = os.path.isfile
store = GCSStore()
# Test without basename
gcs_url = 'gs://bucket/path/to/blob.txt'
store.upload_file(filename=fpath, blob=gcs_url, use_basename=False)
client.return_value.get_bucket.assert_called_with('bucket')
client.return_value.get_bucket.return_value.blob.assert_called_with(
'path/to/blob.txt')
(client.return_value.get_bucket.return_value.blob.return_value.
upload_from_filename.assert_called_with(fpath))
# Test with basename
gcs_url = 'gs://bucket/path/to/'
store.upload_file(filename=fpath, blob=gcs_url, use_basename=True)
client.return_value.get_bucket.assert_called_with('bucket')
client.return_value.get_bucket.return_value.blob.assert_called_with(
'path/to/test.txt')
(client.return_value.get_bucket.return_value.blob.return_value.
upload_from_filename.assert_called_with(fpath))
def test_upload_dir(self, client, _):
dirname1 = tempfile.mkdtemp()
fpath1 = dirname1 + '/test1.txt'
with open(fpath1, 'w') as f:
f.write('data1')
fpath2 = dirname1 + '/test2.txt'
with open(fpath2, 'w') as f:
f.write('data2')
dirname2 = tempfile.mkdtemp(prefix=dirname1 + '/')
fpath3 = dirname2 + '/test3.txt'
with open(fpath3, 'w') as f:
f.write('data3')
(client.return_value
.get_bucket.return_value
.blob.return_value
.upload_from_filename.side_effect) = os.path.isfile
store = GCSStore()
blob_path = 'path/to/'
gcs_url = 'gs://bucket/' + blob_path
rel_path1 = dirname1.split('/')[-1]
rel_path2 = dirname2.split('/')[-1]
# Test without basename
store.upload_dir(dirname=dirname1, blob=gcs_url, use_basename=False)
client.return_value.get_bucket.assert_called_with('bucket')
client.return_value.get_bucket.return_value.blob.assert_has_calls(
[
mock.call('{}test1.txt'.format(blob_path)),
mock.call('{}test2.txt'.format(blob_path)),
mock.call('{}{}/test3.txt'.format(blob_path, rel_path2)),
], any_order=True)
(client.return_value
.get_bucket.return_value
.blob.return_value
.upload_from_filename.assert_has_calls([mock.call(fpath1),
mock.call(fpath2),
mock.call(fpath3),
], any_order=True))
# Test with basename
store.upload_dir(dirname=dirname1, blob=gcs_url, use_basename=True)
client.return_value.get_bucket.assert_called_with('bucket')
client.return_value.get_bucket.return_value.blob.assert_has_calls(
[
mock.call('{}{}/test1.txt'.format(blob_path, rel_path1)),
mock.call('{}{}/test2.txt'.format(blob_path, rel_path1)),
mock.call('{}{}/{}/test3.txt'.format(blob_path, rel_path1, rel_path2)),
], any_order=True)
(client.return_value
.get_bucket.return_value
.blob.return_value
.upload_from_filename.assert_has_calls([mock.call(fpath1),
mock.call(fpath2),
mock.call(fpath3),
], any_order=True))
def get_store(cls, store_type=None, **kwargs):
store_type = store_type or get_from_env(['POLYAXON_STORE'])
# We assume that `None` refers to local store as well
store_type = cls._LOCAL_STORE if store_type is None else store_type
if store_type not in cls._STORE_TYPES:
raise PolyaxonStoresException(
'Received an unrecognised store type `{}`.'.format(store_type))
if store_type == cls._LOCAL_STORE:
from polystores.stores.local_store import LocalStore
return LocalStore()
if store_type == cls._AZURE_STORE:
from polystores.stores.azure_store import AzureStore
return AzureStore(**kwargs)
if store_type == cls._S3_STORE:
from polystores.stores.s3_store import S3Store
return S3Store(**kwargs)
if store_type == cls._GCS_STORE:
from polystores.stores.gcs_store import GCSStore
return GCSStore(**kwargs)
raise PolyaxonStoresException(
'Received an unrecognised store type `{}`.'.format(store_type))
def test_download_dir_with_basename(self, client, _):
dirname1 = tempfile.mkdtemp()
dirname2 = tempfile.mkdtemp(prefix=dirname1 + '/')
def mkfile(fname):
return open(fname, 'w')
(client.return_value.get_bucket.return_value.get_blob.return_value.
download_to_filename.side_effect) = mkfile
store = GCSStore()
blob_path = 'path/to/'
gcs_url = 'gs://bucket/' + blob_path
rel_path2 = dirname2.split('/')[-1]
# Mock return list
obj_mock1 = mock.Mock()
obj_mock1.configure_mock(name=blob_path + 'foo/test1.txt', size=1)
obj_mock2 = mock.Mock()
obj_mock2.configure_mock(name=blob_path + 'foo/test2.txt', size=1)
subdir_mock = mock.Mock()
subdirname = rel_path2 + '/'
subdir_mock.configure_mock(prefixes=(blob_path + 'foo/' +
subdirname, ))
obj_mock3 = mock.Mock()
obj_mock3.configure_mock(name=blob_path + 'foo/' + subdirname +
'test3.txt',
size=1)
mock_results1 = mock.MagicMock()
mock_results1.configure_mock(pages=[subdir_mock])
mock_results1.__iter__.return_value = [obj_mock1, obj_mock2]
mock_results2 = mock.MagicMock()
mock_results2.configure_mock(pages=[])
mock_results2.__iter__.return_value = [obj_mock3]
def list_side_effect(prefix, delimiter='/'):
if prefix == blob_path + 'foo/':
return mock_results1
return mock_results2
client.return_value.get_bucket.return_value.list_blobs.side_effect = list_side_effect
dirname3 = tempfile.mkdtemp()
# Test with basename
store.download_dir(blob=gcs_url + 'foo',
local_path=dirname3,
use_basename=True)
client.return_value.get_bucket.assert_called_with('bucket')
client.return_value.get_bucket().get_blob.assert_has_calls(
[
mock.call('{}foo/test1.txt'.format(blob_path)),
mock.call('{}foo/test2.txt'.format(blob_path)),
mock.call('{}foo/{}/test3.txt'.format(blob_path, rel_path2)),
],
any_order=True)
(client.return_value.get_bucket.return_value.get_blob.return_value.
download_to_filename.assert_has_calls([
mock.call('{}/foo/test1.txt'.format(dirname3)),
mock.call('{}/foo/test2.txt'.format(dirname3)),
mock.call('{}/foo/{}/test3.txt'.format(dirname3, rel_path2))
],
any_order=True))
def test_list_empty(self, client, _):
gcs_url = 'gs://bucket/path/to/blob'
store = GCSStore()
client.return_value.get_bucket.return_value.list_blobs.return_value = mock.MagicMock(
)
assert store.list(gcs_url) == {'blobs': [], 'prefixes': []}
