Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def test_get_gcs_client_with_file_path(self):
mock = MagicMock()
file_path = 'test.json'
os.environ['env_name'] = file_path
with patch('luigi.contrib.gcs.GCSClient'):
with patch('google.oauth2.service_account.Credentials.from_service_account_file', mock):
with patch('os.path.isfile', return_value=True):
GCSConfig(gcs_credential_name='env_name').get_gcs_client()
self.assertEqual(file_path, mock.call_args[0][0])
def test_get_gcs_client_with_json(self):
mock = MagicMock()
json_str = '{"test": 1}'
os.environ['env_name'] = json_str
with patch('luigi.contrib.gcs.GCSClient'):
with patch('google.oauth2.service_account.Credentials.from_service_account_info', mock):
GCSConfig(gcs_credential_name='env_name').get_gcs_client()
self.assertEqual(dict(test=1), mock.call_args[0][0])
def test_get_gcs_client_without_gcs_credential_name(self):
mock = MagicMock()
with patch('luigi.contrib.gcs.GCSClient', mock):
GCSConfig().get_gcs_client()
self.assertEqual(dict(oauth_credentials=None), mock.call_args[1])
def get_timestamp(path: str) -> datetime:
if path.startswith('s3://'):
return S3Config().get_s3_client().get_key(path).last_modified
elif path.startswith('gs://'):
# for gcs object
# should PR to luigi
bucket, obj = GCSConfig().get_gcs_client()._path_to_bucket_and_key(path)
result = GCSConfig().get_gcs_client().client.objects().get(bucket=bucket, object=obj).execute()
return result['updated']
else:
raise
def exists(path: str) -> bool:
if path.startswith('s3://'):
return S3Config().get_s3_client().exists(path)
elif path.startswith('gs://'):
return GCSConfig().get_gcs_client().exists(path)
else:
raise
def __init__(self, workspace_directory: str = '', task_filters: List[str] = [], tqdm_disable: bool = False):
"""must set $GCS_CREDENTIAL"""
self.workspace_directory = workspace_directory
self.task_filters = task_filters
self.tqdm_disable = tqdm_disable
self.gcs_client = GCSConfig().get_gcs_client()
def __init__(self, file_path: str, temporary_directory: str) -> None:
self._file_path = file_path
self._temporary_directory = temporary_directory
self._client = GCSConfig().get_gcs_client()
def get_object_storage_target(path: str, format: Format) -> luigi.Target:
if path.startswith('s3://'):
return luigi.contrib.s3.S3Target(path, client=S3Config().get_s3_client(), format=format)
elif path.startswith('gs://'):
return luigi.contrib.gcs.GCSTarget(path, client=GCSConfig().get_gcs_client(), format=format)
else:
raise