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