def SetUp(self):
"""Called by base class's SetUp() method, which does additional mocking."""
self.SelectApi('v1')
self._mock_compute_client = self._MockClient('compute', 'v1')
self._mock_iam_client = self._MockClient('iam', 'v1')
self._mock_storage_client = self._MockClient('storage', 'v1')
self.diagnose_client = diagnose_utils.DiagnoseClient(
compute_client=self._mock_compute_client,
iam_client=self._mock_iam_client,
storage_client=self._mock_storage_client)
def Run(self, args):
"""See base class."""
self._diagnose_client = diagnose_utils.DiagnoseClient()
instance_ref = self._ResolveInstance(args)
project = properties.VALUES.core.project.Get(required=True)
service_account = self._GetDiagnosticsServiceAccount(project)
expiration_time = self._GetSignedUrlExpiration()
bucket = self._GetLogBucket(project)
log_path = self._GetLogPath(instance_ref.instance)
url = self._CreateResumableSignedUrl(service_account, expiration_time,
bucket, log_path)
diagnostics_entry = self._ConstructDiagnosticsKeyEntry(
url, args.collect_process_traces)
self._diagnose_client.UpdateMetadata(
project, instance_ref, _DIAGNOSTICS_METADATA_KEY, diagnostics_entry)
log.Print(_SUCCESS_MSG.format(bucket, log_path))
return {'bucket': bucket, 'logPath': log_path, 'signedUrl': url}