def setup(self):
self.client = gcsio.GcsIO()
def test_num_retries(self):
# BEAM-7424: update num_retries accordingly if storage_client is
# regenerated.
self.assertEqual(gcsio.GcsIO().client.num_retries, 20)
def test_retry_func(self):
# BEAM-7667: update retry_func accordingly if storage_client is
# regenerated.
self.assertIsNotNone(gcsio.GcsIO().client.retry_func)
def setUp(self):
self.client = FakeGcsClient()
self.gcs = gcsio.GcsIO(self.client)
def gcs_list(gcs_dir):
"""Lists all files in a GCS directory."""
gcs = gcsio.GcsIO()
return gcs.list_prefix(gcs_dir)
def get_file_size(fpath):
fs = FileSystems.get_filesystem(fpath)
if type(fs) == GCSFileSystem:
return gcsio.GcsIO().size(fpath)
else:
return FileIO(fpath, "rb").size()
def open_file_write(fpath):
fs = FileSystems.get_filesystem(fpath)
if type(fs) == GCSFileSystem:
return gcsio.GcsIO().open(fpath, mode='w')
else:
return FileSystems.create(fpath)
def open_file_read(fpath):
fs = FileSystems.get_filesystem(fpath)
if type(fs) == GCSFileSystem:
return gcsio.GcsIO().open(fpath)
else:
return FileSystems.open(fpath)
