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_copy_and_delete_prefix(self):
# Copy a set of objects using the prefix
source_path = "gs://aries_test/test_folder"
dest_path = "gs://aries_test/copy_test/"
objects = gs.GSPrefix(source_path)
try:
objects.copy(dest_path)
copied = StorageFolder(dest_path)
self.assert_object_counts(copied, 1, 1, 3)
self.assertEqual(copied.folder_names[0], "test_folder")
finally:
# Delete the copied files
StorageFolder("gs://aries_test/copy_test/").delete()
def test_intraday_series_with_cache(self):
"""Tests getting daily series data.
"""
data_source = AlphaVantage(self.api_key, cache_folder=self.cache)
df1 = data_source.get_daily_series("AAPL")
self.assert_data_frame(df1)
filenames = StorageFolder(self.cache).file_names
self.assertEqual(len(filenames), 1, "Files: %s" % filenames)
# Get data again, cached data should be returned.
df2 = data_source.get_daily_series("AAPL")
self.assert_data_frame(df1)
self.assertEqual(len(df1), len(df2))
filenames = StorageFolder(self.cache).file_names
self.assertEqual(len(filenames), 1, "Files: %s" % filenames)
def path_exists(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:
logging.debug(f"Checking existence of {path}...")
# Check if path is prefix, and create StoragePrefix object and check if exists
if path.endswith("*"):
return StoragePrefix(path.rstrip("*")).exists()
# Check if it exists as a file or folder, by creating StorageFile and StorageFolder object
return StorageFile(path).exists() or StorageFolder(path).exists()
except RuntimeError as e:
traceback.print_exc()
if str(e) != "":
logging.error(f"StorageHelper error for {job_name}:\n{e}")
return False
except:
traceback.print_exc()
logging.error(f"Unable to check path existence: {path}")
raise
class TestStorageGCP(TestStorage):
SCHEME = "gs"
HOST = "aries_test"
TEST_ROOT_PATH = "/storage_test"
TEST_ROOT = "%s://%s%s" % (SCHEME, HOST, TEST_ROOT_PATH)
test_folder = StorageFolder(TEST_ROOT)
@classmethod
def setUpClass(cls):
cls.CREDENTIALS = False
# Google credentials are required for setting up the class.
gs.setup_credentials(
"GOOGLE_CREDENTIALS",
os.path.join(os.path.dirname(__file__), "gcp.json"))
try:
super().setUpClass()
cls.CREDENTIALS = True
except Exception as ex:
print("%s: %s" % (type(ex), str(ex)))
traceback.print_exc()
def setUp(self):
# Skip test if GCP_ACCESS is not True.
if not self.CREDENTIALS:
self.skipTest("GCP Credentials not found.")
super().setUp()
time.sleep(0.5)
class TestStorageAWS(TestStorage):
SCHEME = "s3"
HOST = "davelab-test"
TEST_ROOT_PATH = "/storage_test"
TEST_ROOT = "%s://%s%s" % (SCHEME, HOST, TEST_ROOT_PATH)
test_folder = StorageFolder(TEST_ROOT)
@classmethod
def setUpClass(cls):
cls.CREDENTIALS = False
# AWS credentials are loaded from environment variable directly.
try:
super().setUpClass()
cls.CREDENTIALS = True
except Exception as ex:
print("%s: %s" % (type(ex), str(ex)))
traceback.print_exc()
raise ex
def setUp(self):
# Skip test if self.CREDENTIALS is not True.
if not self.CREDENTIALS:
self.skipTest("Credentials for %s not found." %
self.__class__.__name__)
super().setUp()
time.sleep(0.5)
def mv(src_path, dest_dir):
# Check if it is remote directory
is_directory = StorageFolder(src_path).exists()
# Convert to Rclone remote structure
src_path = src_path.replace("s3://", "s3:")
dest_dir = dest_dir.replace("s3://", "s3:")
if dest_dir.endswith("/"):
cmd = "copy"
else:
cmd = "copyto"
if src_path.endswith("*"):
basedir, basename = src_path.rsplit("/", 1)
return f"--include {basename} {cmd} {basedir} {dest_dir} --progress"
elif src_path.startswith("s3:"):
if is_directory:
newdir = src_path.rstrip("/").rsplit("/", 1)[-1]
return f"{cmd} {src_path} {dest_dir}/{newdir} --progress"
else:
return f"{cmd} {src_path} {dest_dir} --progress"
else:
newdir = src_path.rstrip("/").rsplit("/", 1)[-1]
return f"{cmd} {src_path} {dest_dir}$([ -d {src_path} ] && echo '/{newdir}') --progress"
class TestWebStorage(AriesTest):
"""Contains test cases for HTTP and FTP.
"""
test_folder_path = os.path.join(os.path.dirname(__file__), "fixtures", "test_folder")
test_folder = StorageFolder(test_folder_path)
@classmethod
def setUpClass(cls):
super().setUpClass()
if not cls.test_folder.exists():
cls.test_folder.create()
def test_http(self):
"""
"""
# URL does not exist
storage_obj = StorageFile("http://example.com/abc/")
self.assertFalse(storage_obj.exists())
# URL exists
storage_obj = StorageFile("https://www.google.com")
self.assertTrue(storage_obj.exists())
# Download. Copy to local file.
storage_obj = StorageFile("https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf")
local_file_path = os.path.join(self.test_folder_path, "test.pdf")
if os.path.exists(local_file_path):
os.remove(local_file_path)
storage_obj.copy(local_file_path)
self.assertTrue(os.path.exists(local_file_path))
self.assertGreater(StorageFile(local_file_path).size, 0)
StorageFile(local_file_path).delete()
def test_get_folder_names(self):
folder = StorageFolder(self.TEST_ROOT)
folder_names = folder.folder_names
# There are two sub-folders
self.assertEqual(len(folder_names), 2, 'Folders: %s' % folder_names)
self.assertIn("test_folder_0", folder_names)
self.assertIn("test_folder_1", folder_names)
# "/" should not be found in folder names
self.assertNotIn("/", str(folder_names))
def test_gs_folder(self):
"""Tests accessing a Google Cloud Storage folder.
"""
# Access a folder in a bucket
self.assert_gs_folder("gs://aries_test/test_folder")
# self.assert_gs_folder("gs://aries_test/test_folder/")
# StorageFolder methods
gs_folder = StorageFolder("gs://aries_test/test_folder/")
self.assertEqual(len(gs_folder.folder_paths), 1)
self.assertEqual(len(gs_folder.file_paths), 1)
def test_parse_uri(self):
"""Tests parsing GCS URI
"""
# File
file_obj = StorageFile("s3://%s/test_file.txt" % self.TEST_BUCKET_NAME)
self.assertEqual(file_obj.scheme, "s3")
self.assertEqual(file_obj.path, "/test_file.txt")
# Folder
folder_obj = StorageFolder("s3://%s/test_folder" %
self.TEST_BUCKET_NAME)
self.assertEqual(folder_obj.uri,
"s3://%s/test_folder/" % self.TEST_BUCKET_NAME)
self.assertEqual(folder_obj.scheme, "s3")
self.assertEqual(folder_obj.path, "/test_folder/")
# Bucket root
folder_obj = StorageFolder("s3://%s" % self.TEST_BUCKET_NAME)
self.assertEqual(folder_obj.uri, "s3://%s/" % self.TEST_BUCKET_NAME)
self.assertEqual(folder_obj.scheme, "s3")
self.assertEqual(folder_obj.path, "/")
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 __intraday_cache_files(self, symbol):
"""Gets all temporary cache data file for intraday data.
Args:
symbol (str): The symbol of the equity/stock.
Returns:
list: A list of StorageFile objects.
"""
prefix = self.__intraday_cache_file_prefix(symbol)
cached_files = []
cache_folder = StorageFolder(self.cache)
for f in cache_folder.files:
if f.basename.startswith(prefix):
cached_files.append(f)
return cached_files
def setUpClass(cls):
gs.setup_credentials("GOOGLE_CREDENTIALS", os.path.join(os.path.dirname(__file__), "gcp.json"))
super().setUpClass()
try:
# Check if GCP is accessible by listing all the buckets
storage.Client().list_buckets(max_results=1)
cls.GCP_ACCESS = True
# Removes test folder if it is already there
StorageFolder("gs://aries_test/copy_test/").delete()
StorageFile("gs://aries_test/copy_test").delete()
StorageFile("gs://aries_test/abc.txt").delete()
StorageFile("gs://aries_test/new_file.txt").delete()
StorageFile("gs://aries_test/moved_file.txt").delete()
StorageFile("gs://aries_test/local_upload.txt").delete()
except Exception as ex:
print("%s: %s" % (type(ex), str(ex)))
traceback.print_exc()
class TestStorageGCP(TestStorageWithCredentials):
SCHEME = "gs"
HOST = "aries_test"
TEST_ROOT_PATH = "/storage_test"
TEST_ROOT = "%s://%s%s" % (SCHEME, HOST, TEST_ROOT_PATH)
test_folder = StorageFolder(TEST_ROOT)
@classmethod
def setUpClass(cls):
try:
# Google credentials are required for setting up the class.
gs.setup_credentials(
"GOOGLE_CREDENTIALS",
os.path.join(os.path.dirname(__file__), "gcp.json"))
super().setUpClass()
cls.CREDENTIALS = True
except Exception as ex:
print("%s: %s" % (type(ex), str(ex)))
traceback.print_exc()
def assert_bucket_root(self, gs_path):
"""Checks if the bucket root contains the expected folder and files.
Args:
gs_path (str): Google cloud storage path to the bucket root, e.g. "gs://bucket_name".
"""
parent = StorageFolder(gs_path)
# Test listing the folders
folders = parent.folders
self.assertEqual(len(folders), 1, str(folders))
self.assertTrue(isinstance(folders[0], StorageFolder), "Type: %s" % type(folders[0]))
self.assertEqual(folders[0].uri, "gs://aries_test/test_folder/")
# Test listing the files
files = parent.files
self.assertEqual(len(files), 2, files)
for file in files:
self.assertTrue(isinstance(file, StorageFile), "Type: %s" % type(file))
self.assertIn(file.uri, [
"gs://aries_test/file_in_root.txt",
"gs://aries_test/test_folder"
])
class TestStorageAWS(TestStorageWithCredentials):
SCHEME = "s3"
HOST = "davelab-test"
TEST_ROOT_PATH = "/storage_test"
TEST_ROOT = "%s://%s%s" % (SCHEME, HOST, TEST_ROOT_PATH)
test_folder = StorageFolder(TEST_ROOT)
@classmethod
def setUpClass(cls):
cls.CREDENTIALS = False
# AWS credentials are loaded from environment variable directly.
# Skip tests if credentials are not configured.
if not os.environ.get("AWS_SECRET_ACCESS_KEY") or not os.environ.get(
"AWS_ACCESS_KEY_ID"):
return
try:
super().setUpClass()
cls.CREDENTIALS = True
except Exception as ex:
print("%s: %s" % (type(ex), str(ex)))
traceback.print_exc()
raise ex
def __init__(self, api_key, cache_folder=None):
"""Initialize the AlphaVantage Data Source
Args:
api_key (str): AlphaVantage API key.
cache_folder (str, optional): Path to local cache data folder. Defaults to None.
CSV file containing the series will be saved into the cache_folder
"""
self.api_key = api_key
self.cache = cache_folder
if self.cache:
self.cache_folder = StorageFolder.init(self.cache)
self.cache_folder.create()
else:
self.cache_folder = None
self.web_api = AlphaVantageAPI(api_key, datatype="csv")
# Expiration time for daily cache data (days)
self.daily_cache_expiration = 1
# Expiration time for intraday cache data (minutes)
self.intraday_cache_expiration = 30
def test_parse_uri(self):
"""Tests parsing GCS URI
"""
# Bucket root without "/"
gs_obj = gs.GSPrefix("gs://aries_test")
self.assertEqual(gs_obj.bucket_name, "aries_test")
self.assertEqual(gs_obj.prefix, "")
# Bucket root with "/"
gs_obj = gs.GSPrefix("gs://aries_test/")
self.assertEqual(gs_obj.bucket_name, "aries_test")
self.assertEqual(gs_obj.prefix, "")
# Object without "/"
gs_obj = gs.GSPrefix("gs://aries_test/test_folder")
self.assertEqual(gs_obj.bucket_name, "aries_test")
self.assertEqual(gs_obj.prefix, "test_folder")
# Object with "/"
gs_obj = gs.GSPrefix("gs://aries_test/test_folder/")
self.assertEqual(gs_obj.bucket_name, "aries_test")
self.assertEqual(gs_obj.prefix, "test_folder/")
# Folder without "/"
gs_obj = StorageFolder("gs://aries_test/test_folder")
self.assertEqual(gs_obj.uri, "gs://aries_test/test_folder/")
self.assertEqual(gs_obj.bucket_name, "aries_test")
self.assertEqual(gs_obj.raw.prefix, "test_folder/")
def create_folder(cls, relative_path):
abs_path = os.path.join(cls.TEST_ROOT, relative_path)
StorageFolder(abs_path).create()
class TestStorage(AriesTest):
SCHEME = "file"
HOST = ""
TEST_ROOT_PATH = os.path.join(os.path.dirname(__file__), "fixtures",
"test_folder")
TEST_ROOT = "%s://%s%s" % (SCHEME, HOST, TEST_ROOT_PATH)
test_folder = StorageFolder(TEST_ROOT)
@classmethod
def create_folder(cls, relative_path):
abs_path = os.path.join(cls.TEST_ROOT, relative_path)
StorageFolder(abs_path).create()
@classmethod
def create_file(cls, relative_path, content):
"""Creates a file relative to the test root
"""
abs_path = os.path.join(cls.TEST_ROOT, relative_path)
with StorageFile.init(abs_path, "w") as f:
f.write(content)
@classmethod
def setUpClass(cls):
super().setUpClass()
if not cls.test_folder.exists():
cls.test_folder.create()
cls.test_folder.empty()
cls.create_folder("test_folder_0")
cls.create_folder("test_folder_1")
cls.create_file("file_in_test_folder", "")
cls.create_file("test_folder_0/empty_file", "")
cls.create_file("test_folder_0/abc.txt", "abc\ncba\n")
cls.create_file("test_folder_1/.hidden_file", "")
@classmethod
def tearDownClass(cls):
StorageFolder(cls.TEST_ROOT).empty()
StorageFolder(cls.TEST_ROOT).delete()
super().tearDownClass()
def setUp(self):
super().setUp()
def test_prefix(self):
storage_prefix = StoragePrefix(
os.path.join(self.TEST_ROOT, "test_folder"))
self.assertEqual(len(storage_prefix.objects), 3)
def test_storage_folder(self):
# Scheme
self.assertEqual(self.test_folder.scheme, self.SCHEME)
# Hostname
if self.test_folder.hostname or self.HOST:
self.assertEqual(self.test_folder.hostname, self.HOST)
# Folder URI will always end with "/"
self.assertTrue(str(self.test_folder).endswith("/"))
# Path
self.assertEqual(self.test_folder.path.rstrip("/"),
self.TEST_ROOT_PATH.rstrip("/"))
# Name
self.assertEqual(self.test_folder.basename,
os.path.basename(self.TEST_ROOT_PATH))
self.assertEqual(self.test_folder.name,
os.path.basename(self.TEST_ROOT_PATH))
# Folder attributes
self.assertGreater(len(self.test_folder.get_folder_attributes()), 1)
# File attributes
self.assertIn("file_in_test_folder",
self.test_folder.get_file_attributes("name"))
def test_get_folder_names(self):
folder = StorageFolder(self.TEST_ROOT)
folder_names = folder.folder_names
# There are two sub-folders
self.assertEqual(len(folder_names), 2, 'Folders: %s' % folder_names)
self.assertIn("test_folder_0", folder_names)
self.assertIn("test_folder_1", folder_names)
# "/" should not be found in folder names
self.assertNotIn("/", str(folder_names))
def test_get_paths(self):
sub_folders = self.test_folder.folder_paths
# There are two sub-folders
self.assertEqual(len(sub_folders), 2, "Folders: %s" % sub_folders)
# TODO: local paths not having scheme?
if self.SCHEME == "file":
self.assertIn(os.path.join(self.TEST_ROOT_PATH, "test_folder_0"),
sub_folders)
self.assertIn(os.path.join(self.TEST_ROOT_PATH, "test_folder_1"),
sub_folders)
else:
# TODO: folders ends with "/"??
self.assertIn(os.path.join(self.TEST_ROOT, "test_folder_0/"),
sub_folders)
self.assertIn(os.path.join(self.TEST_ROOT, "test_folder_1/"),
sub_folders)
def test_get_folder_and_file(self):
# Try to get a folder that does not exist. Should return None
self.assertIsNone(self.test_folder.get_folder("not_exist"))
# Get a sub folder
sub_folder = self.test_folder.get_folder("test_folder_1")
self.assertTrue(
isinstance(sub_folder, StorageFolder),
"Failed to get sub folder. Received %s: %s" %
(sub_folder.__class__, sub_folder))
# Get the filename of hidden_file
filenames = sub_folder.file_names
self.assertEqual(len(filenames), 1)
self.assertEqual(filenames[0], ".hidden_file")
# Get a file that does not exist
f = sub_folder.get_file("not_exist")
self.assertIsNone(f)
# Get the hidden file
f = sub_folder.get_file(".hidden_file")
self.assertTrue(f.exists(), "Failed to get the file.")
self.assertEqual(f.basename, ".hidden_file")
self.assertTrue(
isinstance(f, StorageFile),
"Failed to get storage file from a folder. Received %s: %s" %
(f.__class__, f))
def test_text_read(self):
with StorageFile.init(
os.path.join(self.TEST_ROOT, "file_in_test_folder")) as f:
self.assertEqual(f.size, 0)
self.assertEqual(f.tell(), 0)
self.assertEqual(f.seek(0, 2), 0)
self.assertEqual(len(f.read()), 0)
def test_text_read_write(self):
# Write a new file
temp_file_path = os.path.join(self.TEST_ROOT, "temp_file.txt")
with StorageFile.init(temp_file_path, 'w+') as f:
self.assertTrue(f.writable())
self.assertEqual(f.tell(), 0)
self.assertEqual(f.write("abc"), 3)
self.assertEqual(f.tell(), 3)
f.seek(0)
self.assertEqual(f.read(), "abc")
# TODO: File may not exist on the cloud until it is closed.
# self.assertTrue(f.exists())
f.delete()
def test_binary_read_write(self):
# File does not exist, a new one will be created
file_uri = os.path.join(self.TEST_ROOT, "test.txt")
storage_file = StorageFile(file_uri).open("wb")
self.assertEqual(storage_file.scheme, self.SCHEME)
self.assertTrue(storage_file.seekable())
self.assertFalse(storage_file.readable())
self.assertEqual(storage_file.write(b"abc"), 3)
self.assertEqual(storage_file.tell(), 3)
self.assertEqual(storage_file.write(b"def"), 3)
self.assertEqual(storage_file.tell(), 6)
storage_file.close()
self.assertTrue(storage_file.exists())
storage_file.open('rb')
self.assertEqual(storage_file.read(), b"abcdef")
storage_file.close()
storage_file.delete()
self.assertFalse(storage_file.exists())
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 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 __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
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 tearDownClass(cls):
StorageFolder(cls.TEST_ROOT).empty()
StorageFolder(cls.TEST_ROOT).delete()
super().tearDownClass()
def __get_all_daily_cache(self, symbol):
prefix = self.__daily_cache_prefix(symbol)
logger.debug("Getting cache files with prefix: %s" % prefix)
storage_files = StorageFolder(self.cache).filter_files(prefix)
return storage_files
def test_copy_and_delete_folder(self):
source_path = "gs://aries_test/test_folder/"
# Destination path ends with "/", the original folder name will be preserved.
dest_path = "gs://aries_test/copy_test/"
folder = StorageFolder(source_path)
folder.copy(dest_path)
copied = StorageFolder(dest_path)
self.assert_object_counts(copied, 0, 1, 2)
self.assertEqual(copied.folder_names[0], "test_folder")
# Delete the copied files
copied.delete()
self.assertEqual(len(copied.files), 0)
self.assertEqual(len(copied.folders), 0)
# Copy contents only.
dest_path = "gs://aries_test/copy_test/new_name"
copied = StorageFolder(dest_path)
try:
folder = StorageFolder(source_path)
folder.copy(dest_path, contents_only=True)
logger.debug(copied.objects)
self.assert_object_counts(copied, 1, 1, 2)
self.assertEqual(copied.folder_names[0], "test_subfolder")
finally:
# Delete the copied files
StorageFolder("gs://aries_test/copy_test/").delete()
def test_copy_to_root_and_delete(self):
# Source without "/"
source_path = "gs://aries_test/test_folder/test_subfolder"
# Destination is the bucket root, whether it ends with "/" does not matter.
dest_path = "gs://aries_test"
try:
folder = StorageFolder(source_path)
folder.copy(dest_path)
copied = StorageFolder("gs://aries_test/test_subfolder/")
self.assert_object_counts(copied, 1, 0, 1)
self.assertEqual(copied.file_names[0], "file_in_subfolder.txt")
# Delete the copied files
StorageFolder("gs://aries_test/test_subfolder/").delete()
# With "/"
source_path = "gs://aries_test/test_folder/test_subfolder"
dest_path = "gs://aries_test/"
folder = StorageFolder(source_path)
folder.copy(dest_path)
copied = StorageFolder("gs://aries_test/test_subfolder/")
self.assert_object_counts(copied, 1, 0, 1)
self.assertEqual(copied.file_names[0], "file_in_subfolder.txt")
finally:
# Delete the copied files
StorageFolder("gs://aries_test/test_subfolder/").delete()