def cleanup_image(self):
credentials = self.credentials[self.account]
project = credentials.get('project_id')
object_name = self.status_msg['object_name']
compute_driver = get_gce_compute_driver(credentials)
storage_driver = get_gce_storage_driver(credentials)
self.log_callback.info('Cleaning up image: {0} in region: {1}.'.format(
self.cloud_image_name, self.region))
try:
delete_gce_image(compute_driver, project, self.cloud_image_name)
delete_image_tarball(storage_driver, object_name, self.bucket)
except Exception as error:
msg = 'Failed to cleanup image: {0}'.format(error)
self.log_callback.warning(msg)
self.add_error_msg(msg)
def run_job(self):
self.status = SUCCESS
self.log_callback.info('Uploading image.')
timestamp = None
build_time = self.status_msg.get('build_time', 'unknown')
if self.use_build_time and (build_time != 'unknown'):
timestamp = timestamp_from_epoch(build_time)
elif self.use_build_time and (build_time == 'unknown'):
raise MashUploadException(
'use_build_time set for job but build time is unknown.')
self.cloud_image_name = format_string_with_date(
self.base_cloud_image_name, timestamp=timestamp)
self.request_credentials([self.account])
credentials = self.credentials[self.account]
storage_driver = get_gce_storage_driver(credentials)
object_name = ''.join([self.cloud_image_name, '.tar.gz'])
exists = blob_exists(storage_driver, object_name, self.bucket)
if exists and not self.force_replace_image:
raise MashUploadException(
'Image tarball: {object_name} already exists '
'in bucket: {bucket}. Use force_replace_image '
'to replace the existing tarball.'.format(
object_name=object_name, bucket=self.bucket))
elif exists and self.force_replace_image:
self.log_callback.info(
'Deleting tarball: {0}, in the bucket named: {1}'.format(
object_name, self.bucket))
delete_image_tarball(storage_driver, object_name, self.bucket)
upload_image_tarball(storage_driver, object_name,
self.status_msg['image_file'], self.bucket)
self.status_msg['cloud_image_name'] = self.cloud_image_name
self.status_msg['object_name'] = object_name
self.log_callback.info(
'Uploaded image: {0}, to the bucket named: {1}'.format(
object_name, self.bucket))
def test_get_gce_storage_driver(mock_service_account, mock_storage):
creds = Mock()
mock_service_account.Credentials.from_service_account_info.return_value = creds
get_gce_storage_driver({'project_id': 'project'})
mock_storage.Client.assert_called_once_with('project', creds)