def test_cancel_succeed(self, mock_mlengine_client, mock_kfp_context,
mock_dump_json, mock_display):
returned_job = {'jobId': 'mock_job', 'state': 'SUCCEEDED'}
mock_mlengine_client().get_job.return_value = (returned_job)
wait_job('mock_project', 'mock_job')
cancel_func = mock_kfp_context.call_args[1]['on_cancel']
cancel_func()
mock_mlengine_client().cancel_job.assert_called_with(
'mock_project', 'mock_job')
def test_wait_job_status_fail(self, mock_mlengine_client, mock_kfp_context,
mock_dump_json, mock_display):
returned_job = {
'jobId': 'mock_job',
'trainingInput': {
'modelDir': 'test'
},
'state': 'FAILED'
}
mock_mlengine_client().get_job.return_value = returned_job
with self.assertRaises(RuntimeError):
wait_job('mock_project', 'mock_job')
def test_wait_job_succeed(self, mock_mlengine_client, mock_kfp_context,
mock_dump_json, mock_display):
returned_job = {'jobId': 'mock_job', 'state': 'SUCCEEDED'}
mock_mlengine_client().get_job.return_value = (returned_job)
result = wait_job('mock_project', 'mock_job')
self.assertEqual(returned_job, result)