def test_copy_file_from_gcs(self, mock_default_credentials):
"""Test copying file from GCS works."""
mock_responses = [({
'status': '200',
'content-range': '0-10/11'
}, b'{"test": 1}')]
http_mocks.mock_http_response_sequence(mock_responses)
try:
file_path = file_loader.copy_file_from_gcs('gs://fake/file.json')
with open(file_path, 'rb') as f:
self.assertEqual(b'{"test": 1}', f.read())
finally:
os.unlink(file_path)
def _export_assets(cloudasset_client, config, root_id, content_type):
"""Worker function for exporting assets and downloading dump from GCS.
Args:
cloudasset_client (CloudAssetClient): CloudAsset API client interface.
config (object): Inventory configuration on server.
root_id (str): The name of the parent resource to export assests under.
content_type (ContentTypes): The content type to export.
Returns:
str: The path to the temporary file downloaded from GCS or None on
error.
"""
asset_types = config.get_cai_asset_types()
if not asset_types:
asset_types = DEFAULT_ASSET_TYPES
timeout = config.get_cai_timeout()
timestamp = int(time.time())
export_path = _get_gcs_path(config.get_cai_gcs_path(), content_type,
root_id, timestamp)
try:
LOGGER.info(
'Starting Cloud Asset export for %s under %s to GCS object '
'%s.', content_type, root_id, export_path)
if asset_types:
LOGGER.info('Limiting export to the following asset types: %s',
asset_types)
results = cloudasset_client.export_assets(root_id,
export_path,
content_type=content_type,
asset_types=asset_types,
blocking=True,
timeout=timeout)
LOGGER.debug(
'Cloud Asset export for %s under %s to GCS '
'object %s completed, result: %s.', content_type, root_id,
export_path, results)
except api_errors.ApiExecutionError as e:
LOGGER.warning('API Error getting cloud asset data: %s', e)
return None
except api_errors.OperationTimeoutError as e:
LOGGER.warning('Timeout getting cloud asset data: %s', e)
return None
except ValueError as e:
LOGGER.warning('Invalid root resource id: %s', e)
return None
if 'error' in results:
LOGGER.error(
'Export of cloud asset data had an error, aborting: '
'%s', results)
return None
try:
LOGGER.debug('Downloading Cloud Asset data from GCS to disk.')
return file_loader.copy_file_from_gcs(export_path)
except errors.HttpError as e:
LOGGER.warning('Download of CAI dump from GCS failed: %s', e)
return None