def test_copy_folder(self):
# Source folder to be copied
src_folder_uri = os.path.join(self.TEST_ROOT, "test_folder_0")
sub_folder = StorageFolder(src_folder_uri)
# Source folder should exist
self.assertTrue(sub_folder.exists())
# Destination folder
dst_folder_uri = os.path.join(self.TEST_ROOT, "new_folder",
"test_folder_0")
dst_parent = os.path.join(self.TEST_ROOT, "new_folder")
with TempFolder(dst_parent):
# Destination folder should not exist
dst_folder = StorageFolder(dst_folder_uri)
if dst_folder.exists():
dst_folder.delete()
self.assertFalse(dst_folder.exists())
# Copy the folder
if not dst_parent.endswith("/"):
dst_parent += "/"
logger.debug("Copying from %s into %s" %
(sub_folder.uri, dst_parent))
sub_folder.copy(dst_parent)
# Destination folder should now exist and contain an empty file
self.assertTrue(dst_folder.exists())
file_paths = [f.uri for f in dst_folder.files]
self.assertEqual(len(file_paths), 2)
self.assertIn(os.path.join(dst_folder_uri, "empty_file"),
file_paths)
self.assertIn(os.path.join(dst_folder_uri, "abc.txt"), file_paths)
def get_file_size(self, path, job_name=None, **kwargs):
# Ignore local paths
if self.__get_file_protocol(path) == "Local":
logging.warning(f"Ignoring path '{path}' as it is local on the disk image. Assuming the path is present!")
return True
try:
# Check if path is prefix, and create StoragePrefix object and get its size
if path.endswith("*"):
_size = StoragePrefix(path.rstrip("*")).size
# Check if it path exists as a file or folder, by creating StorageFile and StorageFolder object
else:
_file = StorageFile(path)
_folder = StorageFolder(path)
if _file.exists():
_size = _file.size
elif _folder.exists():
_size = _folder.size
else:
_size = 0
# Convert to GB
return float(_size)/2**30
except BaseException as e:
logging.error(f"Unable to get file size: {path}")
if str(e) != "":
logging.error(f"Received the following msg:\n{e}")
raise
def get_file_size(self, path, job_name=None, **kwargs):
retry_count = kwargs.get("retry_count", 0)
# Ignore local paths
if self.__get_file_protocol(path) == "Local":
logging.warning(f"Ignoring path '{path}' as it is local on the disk image. Assuming the path is present!")
return True
if retry_count < 5:
try:
# Check if path is prefix, and create StoragePrefix object and get its size
if path.endswith("*"):
_size = StoragePrefix(path.rstrip("*")).size
# Check if it path exists as a file or folder, by creating StorageFile and StorageFolder object
else:
_file = StorageFile(path)
_folder = StorageFolder(path)
_size = 0
found = False
trial_count = 0
while not found:
if trial_count > 10:
logging.error(f"Cannot get size of '{path}' as it doesn't exist after multiple trials!")
break
time.sleep(trial_count)
if _file.exists():
_size = _file.size
found = True
elif _folder.exists():
_size = _folder.size
found = True
else:
trial_count += 1
logging.warning(f"Cannot get size of '{path}' as it does not exist! Trial {trial_count}/10")
# Convert to GB
return float(_size)/2**30
except BaseException as e:
logging.error(f"Unable to get file size: {path}")
if str(e) != "":
logging.error(f"Received the following msg:\n{e}")
if "dictionary changed size" in str(e):
kwargs['retry_count'] = retry_count + 1
return self.get_file_size(path, job_name, **kwargs)
raise
else:
logging.warning(f"Failed to get size of '{path}'! Attempted to retrieve size {retry_count + 1} times.")
return 0
def test_create_copy_and_delete_file(self):
new_folder_uri = os.path.join(self.TEST_ROOT, "new_folder")
with TempFolder(new_folder_uri) as folder:
self.assertTrue(folder.is_empty())
# Create a sub folder inside the new folder
sub_folder_uri = os.path.join(new_folder_uri, "sub_folder")
logger.debug(sub_folder_uri)
sub_folder = StorageFolder(sub_folder_uri).create()
self.assertTrue(sub_folder.exists())
# Copy an empty file
src_file_path = os.path.join(self.TEST_ROOT, "test_folder_0",
"empty_file")
dst_file_path = os.path.join(new_folder_uri, "copied_file")
f = StorageFile(src_file_path)
logger.debug(f.exists())
time.sleep(2)
f.copy(dst_file_path)
self.assertTrue(StorageFile(dst_file_path).exists())
# Copy a file with content and replace the empty file
src_file_path = os.path.join(self.TEST_ROOT, "test_folder_0",
"abc.txt")
dst_file_path = os.path.join(new_folder_uri, "copied_file")
f = StorageFile(src_file_path)
f.copy(dst_file_path)
dst_file = StorageFile(dst_file_path)
self.assertTrue(dst_file.exists())
# Use the shortcut to read file, the content will be binary.
self.assertEqual(dst_file.read(), b"abc\ncba\n")
# Empty the folder. This should delete file and sub folder only
folder.empty()
self.assertTrue(folder.exists())
self.assertTrue(folder.is_empty())
self.assertFalse(sub_folder.exists())
self.assertFalse(dst_file.exists())
def assert_gs_folder(self, gs_path):
"""Checks if a Google Cloud Storage folder contains the expected folders and files.
Args:
gs_path ([type]): [description]
"""
# Test listing the folders
parent = StorageFolder(gs_path)
folders = parent.get_folder_attributes()
self.assertTrue(parent.exists())
# self.assertEqual(parent.size, 11)
self.assertEqual(len(folders), 1)
self.assertEqual(folders[0], "gs://aries_test/test_folder/test_subfolder/")
names = parent.folder_names
self.assertEqual(len(names), 1)
self.assertEqual(names[0], "test_subfolder")
# Test listing the files
files = parent.files
self.assertEqual(len(files), 1)
self.assertTrue(isinstance(files[0], StorageFile), "Type: %s" % type(files[0]))
self.assertEqual(files[0].uri, "gs://aries_test/test_folder/file_in_folder.txt")
names = parent.file_names
self.assertEqual(len(names), 1)
self.assertEqual(names[0], "file_in_folder.txt")
def __exit__(self, exc_type, exc_val, exc_tb):
folder = StorageFolder(self.folder_uri)
if folder.exists():
folder.delete()
def __enter__(self):
folder = StorageFolder(self.folder_uri)
if not folder.exists():
folder.create()
return folder
