def test_execute_create_job_fail(self, mock_mlengine_client,
mock_kfp_context, mock_dump_json,
mock_display):
mock_kfp_context().__enter__().context_id.return_value = 'ctx1'
job = {}
mock_mlengine_client().create_job.side_effect = errors.HttpError(
resp=mock.Mock(status=400), content=b'bad request')
with self.assertRaises(errors.HttpError) as context:
create_job('mock_project', job)
self.assertEqual(400, context.exception.resp.status)
def test_create_job_use_context_id_as_name(self, mock_mlengine_client,
mock_kfp_context,
mock_dump_json, mock_display):
context_id = 'ctx1'
job = {}
returned_job = {'jobId': 'job_ctx1', 'state': 'SUCCEEDED'}
mock_mlengine_client().get_job.return_value = (returned_job)
mock_kfp_context().__enter__().context_id.return_value = context_id
create_job('mock_project', job)
mock_mlengine_client().create_job.assert_called_with(
project_id='mock_project', job={'jobId': 'job_ctx1'})
def test_cancel_succeed(self, mock_mlengine_client, mock_kfp_context,
mock_dump_json, mock_display):
mock_kfp_context().__enter__().context_id.return_value = 'ctx1'
job = {}
returned_job = {'jobId': 'job_ctx1', 'state': 'SUCCEEDED'}
mock_mlengine_client().get_job.return_value = (returned_job)
create_job('mock_project', job)
cancel_func = mock_kfp_context.call_args[1]['on_cancel']
cancel_func()
mock_mlengine_client().cancel_job.assert_called_with(
'mock_project', 'job_ctx1')
def test_execute_job_status_fail(self, mock_mlengine_client,
mock_kfp_context, mock_dump_json,
mock_display):
mock_kfp_context().__enter__().context_id.return_value = 'ctx1'
job = {}
returned_job = {
'jobId': 'mock_job',
'trainingInput': {
'modelDir': 'test'
},
'state': 'FAILED'
}
mock_mlengine_client().get_job.return_value = returned_job
with self.assertRaises(RuntimeError):
create_job('mock_project', job)
def test_execute_conflict_fail(self, mock_mlengine_client,
mock_kfp_context, mock_dump_json,
mock_display):
mock_kfp_context().__enter__().context_id.return_value = 'ctx1'
job = {}
returned_job = {
'jobId': 'job_ctx1',
'trainingInput': {
'modelDir': 'test'
},
'state': 'SUCCEEDED'
}
mock_mlengine_client().create_job.side_effect = errors.HttpError(
resp=mock.Mock(status=409), content=b'conflict')
mock_mlengine_client().get_job.return_value = returned_job
with self.assertRaises(errors.HttpError) as context:
create_job('mock_project', job)
self.assertEqual(409, context.exception.resp.status)
def test_create_job_succeed(self, mock_mlengine_client, mock_kfp_context,
mock_dump_json, mock_display):
mock_kfp_context().__enter__().context_id.return_value = 'ctx1'
job = {}
returned_job = {'jobId': 'job_ctx1', 'state': 'SUCCEEDED'}
mock_mlengine_client().get_job.return_value = (returned_job)
result = create_job('mock_project', job)
self.assertEqual(returned_job, result)
mock_mlengine_client().create_job.assert_called_with(
project_id='mock_project', job={'jobId': 'job_ctx1'})
def test_execute_retry_job_success(self, mock_mlengine_client,
mock_kfp_context, mock_dump_json,
mock_display):
mock_kfp_context().__enter__().context_id.return_value = 'ctx1'
job = {}
returned_job = {'jobId': 'job_ctx1', 'state': 'SUCCEEDED'}
mock_mlengine_client().create_job.side_effect = errors.HttpError(
resp=mock.Mock(status=409), content=b'conflict')
mock_mlengine_client().get_job.return_value = returned_job
result = create_job('mock_project', job)
self.assertEqual(returned_job, result)
