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)