def test_launch_template_retry_succeed(self, mock_client, mock_context,
mock_display):
mock_context().__enter__().context_id.return_value = 'ctx-1'
# The job with same name already exists.
mock_client().list_aggregated_jobs.return_value = {
'jobs': [{
'id': 'job-1',
'name': 'test_job-ctx-1'
}]
}
pending_job = {'currentState': 'JOB_STATE_PENDING'}
expected_job = {'id': 'job-1', 'currentState': 'JOB_STATE_DONE'}
mock_client().get_job.side_effect = [pending_job, expected_job]
result = launch_template('project-1',
'gs://foo/bar', {
"parameters": {
"foo": "bar"
},
"environment": {
"zone": "us-central1"
}
},
job_name_prefix='test-job',
wait_interval=0)
self.assertEqual(expected_job, result)
mock_client().launch_template.assert_not_called()
def test_launch_template_retry_succeed(self,
mock_client, mock_context, mock_display, mock_storage):
mock_context().__enter__().context_id.return_value = 'ctx-1'
mock_storage.Client().bucket().blob().exists.return_value = True
mock_storage.Client().bucket().blob().download_as_string.return_value = b'job-1,'
pending_job = {
'currentState': 'JOB_STATE_PENDING'
}
expected_job = {
'id': 'job-1',
'currentState': 'JOB_STATE_DONE'
}
mock_client().get_job.side_effect = [pending_job, expected_job]
result = launch_template('project-1', 'gs://foo/bar', {
"parameters": {
"foo": "bar"
},
"environment": {
"zone": "us-central1"
}
}, staging_dir='gs://staging/dir', wait_interval=0)
self.assertEqual(expected_job, result)
mock_client().launch_template.assert_not_called()
def test_launch_template_succeed(self, mock_client, mock_context, mock_display,
mock_storage):
mock_context().__enter__().context_id.return_value = 'context-1'
mock_storage.Client().bucket().blob().exists.return_value = False
mock_client().launch_template.return_value = {
'job': { 'id': 'job-1' }
}
expected_job = {
'id': 'job-1',
'currentState': 'JOB_STATE_DONE'
}
mock_client().get_job.return_value = expected_job
result = launch_template('project-1', 'gs://foo/bar', {
"parameters": {
"foo": "bar"
},
"environment": {
"zone": "us-central1"
}
}, staging_dir='gs://staging/dir')
self.assertEqual(expected_job, result)
mock_client().launch_template.assert_called_once()
mock_storage.Client().bucket().blob().upload_from_string.assert_called_with(
'job-1,'
)
def test_launch_template_fail(self, mock_client, mock_context,
mock_display):
mock_context().__enter__().context_id.return_value = 'context-1'
mock_client().list_aggregated_jobs.return_value = {'jobs': []}
mock_client().launch_template.return_value = {'job': {'id': 'job-1'}}
failed_job = {'id': 'job-1', 'currentState': 'JOB_STATE_FAILED'}
mock_client().get_job.return_value = failed_job
self.assertRaises(
RuntimeError, lambda: launch_template(
'project-1', 'gs://foo/bar', {
"parameters": {
"foo": "bar"
},
"environment": {
"zone": "us-central1"
}
}))
def test_launch_template_succeed(self, mock_client, mock_context,
mock_display):
mock_context().__enter__().context_id.return_value = 'context-1'
mock_client().list_aggregated_jobs.return_value = {'jobs': []}
mock_client().launch_template.return_value = {'job': {'id': 'job-1'}}
expected_job = {'id': 'job-1', 'currentState': 'JOB_STATE_DONE'}
mock_client().get_job.return_value = expected_job
result = launch_template('project-1', 'gs://foo/bar', {
"parameters": {
"foo": "bar"
},
"environment": {
"zone": "us-central1"
}
})
self.assertEqual(expected_job, result)
mock_client().launch_template.assert_called_once()
