def testSuccessCreateTrainingJobWithOptionalArgs(self):
training_input = copy.deepcopy(self.TRAINING_INPUT)
training_input['trainingInput']['runtimeVersion'] = '1.6'
training_input['trainingInput']['pythonVersion'] = '3.5'
training_input['trainingInput'][
'jobDir'] = 'gs://some-bucket/jobs/test_training'
with patch('airflow.contrib.operators.mlengine_operator.MLEngineHook') \
as mock_hook:
success_response = self.TRAINING_INPUT.copy()
success_response['state'] = 'SUCCEEDED'
hook_instance = mock_hook.return_value
hook_instance.create_job.return_value = success_response
training_op = MLEngineTrainingOperator(
runtime_version='1.6',
python_version='3.5',
job_dir='gs://some-bucket/jobs/test_training',
**self.TRAINING_DEFAULT_ARGS)
training_op.execute(None)
mock_hook.assert_called_with(gcp_conn_id='google_cloud_default',
delegate_to=None)
# Make sure only 'create_job' is invoked on hook instance
self.assertEquals(len(hook_instance.mock_calls), 1)
hook_instance.create_job.assert_called_with(
'test-project', training_input, ANY)
def testSuccessCreateTrainingJobWithOptionalArgs(self):
training_input = copy.deepcopy(self.TRAINING_INPUT)
training_input['trainingInput']['runtimeVersion'] = '1.6'
training_input['trainingInput']['pythonVersion'] = '3.5'
training_input['trainingInput']['jobDir'] = 'gs://some-bucket/jobs/test_training'
with patch('airflow.contrib.operators.mlengine_operator.MLEngineHook') \
as mock_hook:
success_response = self.TRAINING_INPUT.copy()
success_response['state'] = 'SUCCEEDED'
hook_instance = mock_hook.return_value
hook_instance.create_job.return_value = success_response
training_op = MLEngineTrainingOperator(
runtime_version='1.6',
python_version='3.5',
job_dir='gs://some-bucket/jobs/test_training',
**self.TRAINING_DEFAULT_ARGS)
training_op.execute(None)
mock_hook.assert_called_with(gcp_conn_id='google_cloud_default',
delegate_to=None)
# Make sure only 'create_job' is invoked on hook instance
self.assertEqual(len(hook_instance.mock_calls), 1)
hook_instance.create_job.assert_called_with(
'test-project', training_input, ANY)
def testSuccessCreateTrainingJob(self):
with patch('airflow.contrib.operators.mlengine_operator.MLEngineHook') \
as mock_hook:
success_response = self.TRAINING_INPUT.copy()
success_response['state'] = 'SUCCEEDED'
hook_instance = mock_hook.return_value
hook_instance.create_job.return_value = success_response
training_op = MLEngineTrainingOperator(**self.TRAINING_DEFAULT_ARGS)
training_op.execute(None)
mock_hook.assert_called_with(gcp_conn_id='google_cloud_default',
delegate_to=None)
# Make sure only 'create_job' is invoked on hook instance
self.assertEquals(len(hook_instance.mock_calls), 1)
hook_instance.create_job.assert_called_with(
'test-project', self.TRAINING_INPUT, ANY)
def testSuccessCreateTrainingJob(self):
with patch('airflow.contrib.operators.mlengine_operator.MLEngineHook') \
as mock_hook:
success_response = self.TRAINING_INPUT.copy()
success_response['state'] = 'SUCCEEDED'
hook_instance = mock_hook.return_value
hook_instance.create_job.return_value = success_response
training_op = MLEngineTrainingOperator(
**self.TRAINING_DEFAULT_ARGS)
training_op.execute(None)
mock_hook.assert_called_with(gcp_conn_id='google_cloud_default',
delegate_to=None)
# Make sure only 'create_job' is invoked on hook instance
self.assertEquals(len(hook_instance.mock_calls), 1)
hook_instance.create_job.assert_called_with(
'test-project', self.TRAINING_INPUT, ANY)
def testFailedJobError(self):
with patch('airflow.contrib.operators.mlengine_operator.MLEngineHook') \
as mock_hook:
failure_response = self.TRAINING_INPUT.copy()
failure_response['state'] = 'FAILED'
failure_response['errorMessage'] = 'A failure message'
hook_instance = mock_hook.return_value
hook_instance.create_job.return_value = failure_response
with self.assertRaises(RuntimeError) as context:
training_op = MLEngineTrainingOperator(
**self.TRAINING_DEFAULT_ARGS)
training_op.execute(None)
mock_hook.assert_called_with(
gcp_conn_id='google_cloud_default', delegate_to=None)
# Make sure only 'create_job' is invoked on hook instance
self.assertEqual(len(hook_instance.mock_calls), 1)
hook_instance.create_job.assert_called_with(
'test-project', self.TRAINING_INPUT, ANY)
self.assertEqual('A failure message', str(context.exception))
def testFailedJobError(self):
with patch('airflow.contrib.operators.mlengine_operator.MLEngineHook') \
as mock_hook:
failure_response = self.TRAINING_INPUT.copy()
failure_response['state'] = 'FAILED'
failure_response['errorMessage'] = 'A failure message'
hook_instance = mock_hook.return_value
hook_instance.create_job.return_value = failure_response
with self.assertRaises(RuntimeError) as context:
training_op = MLEngineTrainingOperator(
**self.TRAINING_DEFAULT_ARGS)
training_op.execute(None)
mock_hook.assert_called_with(gcp_conn_id='google_cloud_default',
delegate_to=None)
# Make sure only 'create_job' is invoked on hook instance
self.assertEquals(len(hook_instance.mock_calls), 1)
hook_instance.create_job.assert_called_with(
'test-project', self.TRAINING_INPUT, ANY)
self.assertEquals('A failure message', str(context.exception))
def testHttpError(self):
http_error_code = 403
with patch('airflow.contrib.operators.mlengine_operator.MLEngineHook') \
as mock_hook:
hook_instance = mock_hook.return_value
hook_instance.create_job.side_effect = errors.HttpError(
resp=httplib2.Response({'status': http_error_code}),
content=b'Forbidden')
with self.assertRaises(errors.HttpError) as context:
training_op = MLEngineTrainingOperator(
**self.TRAINING_DEFAULT_ARGS)
training_op.execute(None)
mock_hook.assert_called_with(gcp_conn_id='google_cloud_default',
delegate_to=None)
# Make sure only 'create_job' is invoked on hook instance
self.assertEquals(len(hook_instance.mock_calls), 1)
hook_instance.create_job.assert_called_with(
'test-project', self.TRAINING_INPUT, ANY)
self.assertEquals(http_error_code, context.exception.resp.status)
def testHttpError(self):
http_error_code = 403
with patch('airflow.contrib.operators.mlengine_operator.MLEngineHook') \
as mock_hook:
hook_instance = mock_hook.return_value
hook_instance.create_job.side_effect = errors.HttpError(
resp=httplib2.Response({
'status': http_error_code
}), content=b'Forbidden')
with self.assertRaises(errors.HttpError) as context:
training_op = MLEngineTrainingOperator(
**self.TRAINING_DEFAULT_ARGS)
training_op.execute(None)
mock_hook.assert_called_with(
gcp_conn_id='google_cloud_default', delegate_to=None)
# Make sure only 'create_job' is invoked on hook instance
self.assertEquals(len(hook_instance.mock_calls), 1)
hook_instance.create_job.assert_called_with(
'test-project', self.TRAINING_INPUT, ANY)
self.assertEquals(http_error_code, context.exception.resp.status)