def test_list_non_empty(self, client):
base_path = "/path/"
# Create some files to return
dir_prefix = BlobPrefix()
dir_prefix.name = base_path + "dir"
blob1_props = BlobProperties()
blob1_props.size = 42
blob1_props.name = base_path + "file1"
blob2_props = BlobProperties()
blob2_props.size = 0
blob2_props.name = base_path + "file2"
client.get_container_client().walk_blobs.return_value = MockBlobList(
[dir_prefix, blob1_props, blob2_props]
)
store = AzureBlobStoreService()
store.set_connection(connection=client)
key_path = self.wasbs_base + base_path
results = store.list(key=key_path)
assert len(results["blobs"]) == 2
assert len(results["prefixes"]) == 1
assert results["prefixes"][0] == "dir"
assert results["blobs"][0][0] == "file1"
assert results["blobs"][0][1] == 42
assert results["blobs"][1][0] == "file2"
assert results["blobs"][1][1] == 0
def test_parse_wasbs_url(self):
# Correct url
wasbs_url = "wasbs://[email protected]/path"
parsed_url = AzureBlobStoreService.parse_wasbs_url(wasbs_url)
assert parsed_url == ("container", "user", "path")
wasbs_url = "wasbs://[email protected]/"
parsed_url = AzureBlobStoreService.parse_wasbs_url(wasbs_url)
assert parsed_url == ("container", "user", "")
wasbs_url = "wasbs://[email protected]"
parsed_url = AzureBlobStoreService.parse_wasbs_url(wasbs_url)
assert parsed_url == ("container", "user", "")
wasbs_url = "wasbs://[email protected]/path/to/file"
parsed_url = AzureBlobStoreService.parse_wasbs_url(wasbs_url)
assert parsed_url == ("container", "user", "path/to/file")
# Wrong url
wasbs_url = "wasbs://[email protected]/path/to/file"
with self.assertRaises(PolyaxonStoresException):
AzureBlobStoreService.parse_wasbs_url(wasbs_url)
wasbs_url = "wasbs://[email protected]/path/to/file"
with self.assertRaises(PolyaxonStoresException):
AzureBlobStoreService.parse_wasbs_url(wasbs_url)
wasbs_url = "wasbs://[email protected]/path/to/file"
with self.assertRaises(PolyaxonStoresException):
AzureBlobStoreService.parse_wasbs_url(wasbs_url)
def test_list_empty(self, client):
client.get_container_client().walk_blobs.return_value = MockBlobList([])
store = AzureBlobStoreService()
store.set_connection(connection=client)
key_path = self.wasbs_base + "path"
assert store.list(key=key_path) == {"blobs": [], "prefixes": []}
def test_delete_file(self, client):
client.return_value.list_blobs.return_value = MockBlobList([])
base_path = "path/test.txt"
store = AzureBlobStoreService()
key_path = self.wasbs_base + base_path
# Test without basename
store.delete(key_path)
client.return_value.delete_blob.assert_called_with("container", "path/test.txt")
def test_download_file(self, client):
client.return_value.walk_blobs.return_value = MockBlobList([])
dirname = tempfile.mkdtemp()
fpath = dirname + "/test.txt"
def mkfile(buffer):
return open(buffer.name, "w")
client.get_container_client().download_blob().readinto.side_effect = mkfile
base_path = "path/test.txt"
store = AzureBlobStoreService()
store.set_connection(connection=client)
key_path = self.wasbs_base + base_path
# Test without basename
store.download_file(key_path, fpath, use_basename=False)
client.get_container_client().download_blob.assert_called_with(base_path)
assert os.path.exists(fpath)
# Test without basename
store.download_file(key_path, dirname, use_basename=True)
client.get_container_client().download_blob.assert_called_with(base_path)
assert os.path.exists(fpath)
def test_upload_file(self, client):
dirname = tempfile.mkdtemp()
fpath = dirname + "/test.txt"
open(fpath, "w")
base_path = "path/"
base_path_file = base_path + "test.txt"
store = AzureBlobStoreService()
store.set_connection(connection=client)
# Test without basename
key_path = self.wasbs_base + base_path_file
store.upload_file(filename=fpath, blob=key_path, use_basename=False)
call_arg1, call_arg2 = client.get_container_client(
).upload_blob.call_args[0]
assert call_arg1 == base_path_file
assert call_arg2.name == fpath
# Test with basename
key_path = self.wasbs_base + base_path
store.upload_file(filename=fpath, blob=key_path, use_basename=True)
call_arg1, call_arg2 = client.get_container_client(
).upload_blob.call_args[0]
assert call_arg1 == base_path_file
assert call_arg2.name == fpath
def test_upload_dir_with_last_time(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")
last_time = to_datetime(os.stat(fpath2).st_mtime)
time.sleep(0.1)
dirname2 = tempfile.mkdtemp(prefix=dirname1 + "/")
fpath3 = dirname2 + "/test3.txt"
with open(fpath3, "w") as f:
f.write("data3")
store = AzureBlobStoreService()
store.set_connection(connection=client)
blob_path = "path/to/"
azure_url = self.wasbs_base + blob_path
rel_path1 = dirname1.split("/")[-1]
rel_path2 = dirname2.split("/")[-1]
# Test without basename
store.upload_dir(dirname=dirname1,
blob=azure_url,
use_basename=False,
last_time=last_time)
call_args_list = client.get_container_client(
).upload_blob.call_args_list
for call_args in call_args_list:
call_arg1, call_arg2 = call_args[0]
if call_arg1 == "{}{}/test3.txt".format(blob_path, rel_path2):
assert call_arg2.name == fpath3
else:
assert False
# Test with basename
store.upload_dir(dirname=dirname1,
blob=azure_url,
use_basename=True,
last_time=last_time)
call_args_list = client.get_container_client(
).upload_blob.call_args_list[1:]
for call_args in call_args_list:
call_arg1, call_arg2 = call_args[0]
if call_arg1 == "{}{}/{}/test3.txt".format(blob_path, rel_path1,
rel_path2):
assert call_arg2.name == fpath3
else:
assert False
def test_list_non_empty(self, client):
base_path = "/path/"
# Create some files to return
dir_prefix = BlobPrefix()
dir_prefix.name = base_path + "dir"
blob_props = BlobProperties()
blob_props.content_length = 42
blob = Blob(base_path + "file", props=blob_props)
client.return_value.list_blobs.return_value = MockBlobList([dir_prefix, blob])
store = AzureBlobStoreService()
key_path = self.wasbs_base + base_path
results = store.list(key=key_path)
assert len(results["blobs"]) == 1
assert len(results["prefixes"]) == 1
assert results["prefixes"][0] == "dir"
assert results["blobs"][0][0] == "file"
assert results["blobs"][0][1] == 42
def test_download_file(self, client):
client.return_value.list_blobs.return_value = MockBlobList([])
dirname = tempfile.mkdtemp()
fpath = dirname + "/test.txt"
def mkfile(container, cloud_path, fname):
return open(fname, "w")
client.return_value.get_blob_to_path.side_effect = mkfile
base_path = "path/test.txt"
store = AzureBlobStoreService()
key_path = self.wasbs_base + base_path
# Test without basename
store.download_file(key_path, fpath, use_basename=False)
client.return_value.get_blob_to_path.assert_called_with(
"container", base_path, fpath
)
# Test without basename
store.download_file(key_path, dirname, use_basename=True)
client.return_value.get_blob_to_path.assert_called_with(
"container", base_path, fpath
)
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_blob_to_path.side_effect = mkfile
store = AzureBlobStoreService()
store.set_connection(connection=client)
blob_path = "/path/to/"
azure_url = self.wasbs_base + blob_path
rel_path2 = dirname2.split("/")[-1]
# Mock return list
obj_mock1 = BlobProperties()
obj_mock1.size = 42
obj_mock1.name = blob_path + "foo/test1.txt"
obj_mock2 = BlobProperties()
obj_mock2.size = 42
obj_mock2.name = blob_path + "foo/test2.txt"
subdir_mock = BlobPrefix()
subdir_mock.name = blob_path + "foo/" + rel_path2
obj_mock3 = BlobProperties()
obj_mock3.size = 42
obj_mock3.name = blob_path + "foo/" + rel_path2 + "/" + "test3.txt"
# Create some files to return
def list_side_effect(container_name, prefix, delimiter="/", marker=None):
if prefix == blob_path + "foo/":
return MockBlobList([subdir_mock, obj_mock1, obj_mock2])
return MockBlobList([obj_mock3])
client.return_value.walk_blobs.side_effect = list_side_effect
dirname3 = tempfile.mkdtemp()
# Test without basename
store.download_dir(
blob=azure_url + "foo", local_path=dirname3, use_basename=True
)
call_args_list = client.get_container_client().download_blob.call_args_list
for call_args in call_args_list:
call_arg1 = call_args[0]
if call_arg1 == "{}foo/test1.txt".format(blob_path):
assert os.path.exists("{}/foo/test1.txt".format(dirname3))
elif call_arg1 == "{}foo/test2.txt".format(blob_path):
assert os.path.exists("{}/foo/test2.txt".format(dirname3))
elif call_arg1 == "{}foo/{}/test3.txt".format(blob_path, rel_path2):
assert os.path.exists("{}/foo/{}/test3.txt".format(dirname3, rel_path2))
else:
assert False
def test_delete_file(self, client):
client.return_value.walk_blobs.return_value = MockBlobList([])
base_path = "path/test.txt"
store = AzureBlobStoreService()
store.set_connection(connection=client)
key_path = self.wasbs_base + base_path
# Test without basename
store.delete(key_path)
client.get_container_client().delete_blob.assert_called_with("path/test.txt")
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")
store = AzureBlobStoreService()
blob_path = "path/to/"
azure_url = self.wasbs_base + blob_path
rel_path1 = dirname1.split("/")[-1]
rel_path2 = dirname2.split("/")[-1]
# Test without basename
store.upload_dir(dirname=dirname1, blob=azure_url, use_basename=False)
client.return_value.create_blob_from_path.assert_has_calls(
[
mock.call("container", "{}test1.txt".format(blob_path), fpath1),
mock.call("container", "{}test2.txt".format(blob_path), fpath2),
mock.call(
"container", "{}{}/test3.txt".format(blob_path, rel_path2), fpath3
),
],
any_order=True,
)
# Test with basename
store.upload_dir(dirname=dirname1, blob=azure_url, use_basename=True)
client.return_value.create_blob_from_path.assert_has_calls(
[
mock.call(
"container", "{}{}/test1.txt".format(blob_path, rel_path1), fpath1
),
mock.call(
"container", "{}{}/test2.txt".format(blob_path, rel_path1), fpath2
),
mock.call(
"container",
"{}{}/{}/test3.txt".format(blob_path, rel_path1, rel_path2),
fpath3,
),
],
any_order=True,
)
def get_connection_from_type(connection_type: V1ConnectionType, **kwargs):
# We assume that `None` refers to local store as well
if not connection_type or connection_type.kind in {
V1ConnectionKind.VOLUME_CLAIM,
V1ConnectionKind.HOST_PATH,
}:
from polyaxon.stores.local_store import LocalStore
return LocalStore()
if connection_type.kind == V1ConnectionKind.WASB:
from polyaxon.connections.azure.azure_blobstore import AzureBlobStoreService
return AzureBlobStoreService(connection_name=connection_type.name,
**kwargs)
if connection_type.kind == V1ConnectionKind.S3:
from polyaxon.connections.aws.s3 import S3Service
return S3Service(connection_name=connection_type.name, **kwargs)
if connection_type.kind == V1ConnectionKind.GCS:
from polyaxon.connections.gcp.gcs import GCSService
return GCSService(connection_name=connection_type.name, **kwargs)
def test_upload_file(self, client):
dirname = tempfile.mkdtemp()
fpath = dirname + "/test.txt"
open(fpath, "w")
base_path = "path/"
base_path_file = base_path + "test.txt"
store = AzureBlobStoreService()
# Test without basename
key_path = self.wasbs_base + base_path_file
store.upload_file(filename=fpath, blob=key_path, use_basename=False)
client.return_value.create_blob_from_path.assert_called_with(
"container", base_path_file, fpath
)
# Test with basename
key_path = self.wasbs_base + base_path
store.upload_file(filename=fpath, blob=key_path, use_basename=True)
client.return_value.create_blob_from_path.assert_called_with(
"container", base_path_file, fpath
)
def test_list_empty(self, client):
client.return_value.list_blobs.return_value = MockBlobList([])
store = AzureBlobStoreService()
key_path = self.wasbs_base + "path"
assert store.list(key=key_path) == {"blobs": [], "prefixes": []}
def test_download_dir_with_basename(self, client):
dirname1 = tempfile.mkdtemp()
dirname2 = tempfile.mkdtemp(prefix=dirname1 + "/")
def mkfile(container, cloud_path, fname):
return open(fname, "w")
client.return_value.get_blob_to_path.side_effect = mkfile
store = AzureBlobStoreService()
blob_path = "/path/to/"
azure_url = self.wasbs_base + blob_path
rel_path2 = dirname2.split("/")[-1]
# Mock return list
blob_props = BlobProperties()
blob_props.content_length = 42
obj_mock1 = Blob(blob_path + "foo/test1.txt", props=blob_props)
blob_props = BlobProperties()
blob_props.content_length = 42
obj_mock2 = Blob(blob_path + "foo/test2.txt", props=blob_props)
subdir_mock = BlobPrefix()
subdir_mock.name = blob_path + "foo/" + rel_path2
blob_props = BlobProperties()
blob_props.content_length = 42
obj_mock3 = Blob(
blob_path + "foo/" + rel_path2 + "/" + "test3.txt", props=blob_props
)
# Create some files to return
def list_side_effect(container_name, prefix, delimiter="/", marker=None):
if prefix == blob_path + "foo/":
return MockBlobList([subdir_mock, obj_mock1, obj_mock2])
return MockBlobList([obj_mock3])
client.return_value.list_blobs.side_effect = list_side_effect
dirname3 = tempfile.mkdtemp()
# Test without basename
store.download_dir(
blob=azure_url + "foo", local_path=dirname3, use_basename=True
)
client.return_value.get_blob_to_path.assert_has_calls(
[
mock.call(
"container",
"{}foo/test1.txt".format(blob_path),
"{}/foo/test1.txt".format(dirname3),
),
mock.call(
"container",
"{}foo/test2.txt".format(blob_path),
"{}/foo/test2.txt".format(dirname3),
),
mock.call(
"container",
"{}foo/{}/test3.txt".format(blob_path, rel_path2),
"{}/foo/{}/test3.txt".format(dirname3, rel_path2),
),
],
any_order=True,
)
