def test_success_cancel_training_job(self, mock_hook):
success_response = {}
hook_instance = mock_hook.return_value
hook_instance.cancel_job.return_value = success_response
cancel_training_op = MLEngineTrainingCancelJobOperator(
**self.TRAINING_DEFAULT_ARGS)
cancel_training_op.execute(None)
mock_hook.assert_called_once_with(
gcp_conn_id='google_cloud_default', delegate_to=None)
# Make sure only 'cancel_job' is invoked on hook instance
self.assertEqual(len(hook_instance.mock_calls), 1)
hook_instance.cancel_job.assert_called_once_with(
project_id=self.TRAINING_DEFAULT_ARGS['project_id'], job_id=self.TRAINING_DEFAULT_ARGS['job_id'])
def test_http_error(self, mock_hook):
http_error_code = 403
hook_instance = mock_hook.return_value
hook_instance.cancel_job.side_effect = HttpError(
resp=httplib2.Response({'status': http_error_code}), content=b'Forbidden'
)
with self.assertRaises(HttpError) as context:
cancel_training_op = MLEngineTrainingCancelJobOperator(**self.TRAINING_DEFAULT_ARGS)
cancel_training_op.execute(None)
mock_hook.assert_called_once_with(
gcp_conn_id='google_cloud_default', delegate_to=None, impersonation_chain=None,
)
# Make sure only 'create_job' is invoked on hook instance
self.assertEqual(len(hook_instance.mock_calls), 1)
hook_instance.cancel_job.assert_called_once_with(
project_id=self.TRAINING_DEFAULT_ARGS['project_id'], job_id=self.TRAINING_DEFAULT_ARGS['job_id']
)
self.assertEqual(http_error_code, context.exception.resp.status)