def testHttpError(self): http_error_code = 403 with patch('airflow.contrib.operators.mlengine_operator.MLEngineHook') \ as mock_hook: input_with_model = self.INPUT_MISSING_ORIGIN.copy() input_with_model['modelName'] = \ 'projects/experimental/models/test_model' 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: prediction_task = MLEngineBatchPredictionOperator( job_id='test_prediction', project_id='test-project', region=input_with_model['region'], data_format=input_with_model['dataFormat'], input_paths=input_with_model['inputPaths'], output_path=input_with_model['outputPath'], model_name=input_with_model['modelName'].split('/')[-1], dag=self.dag, task_id='test-prediction') prediction_task.execute(None) mock_hook.assert_called_with('google_cloud_default', None) hook_instance.create_job.assert_called_with( 'test-project', { 'jobId': 'test_prediction', 'predictionInput': input_with_model }, ANY) self.assertEquals(http_error_code, context.exception.resp.status)
def testSuccessWithModel(self): with patch('airflow.contrib.operators.mlengine_operator.MLEngineHook') \ as mock_hook: input_with_model = self.INPUT_MISSING_ORIGIN.copy() input_with_model['modelName'] = \ 'projects/test-project/models/test_model' success_message = self.SUCCESS_MESSAGE_MISSING_INPUT.copy() success_message['predictionInput'] = input_with_model hook_instance = mock_hook.return_value hook_instance.get_job.side_effect = errors.HttpError( resp=httplib2.Response({'status': 404}), content=b'some bytes') hook_instance.create_job.return_value = success_message prediction_task = MLEngineBatchPredictionOperator( job_id='test_prediction', project_id='test-project', region=input_with_model['region'], data_format=input_with_model['dataFormat'], input_paths=input_with_model['inputPaths'], output_path=input_with_model['outputPath'], model_name=input_with_model['modelName'].split('/')[-1], dag=self.dag, task_id='test-prediction') prediction_output = prediction_task.execute(None) mock_hook.assert_called_with('google_cloud_default', None) hook_instance.create_job.assert_called_once_with( 'test-project', { 'jobId': 'test_prediction', 'predictionInput': input_with_model }, ANY) self.assertEquals(success_message['predictionOutput'], prediction_output)
def testHttpError(self): http_error_code = 403 with patch('airflow.contrib.operators.mlengine_operator.MLEngineHook') \ as mock_hook: input_with_model = self.INPUT_MISSING_ORIGIN.copy() input_with_model['modelName'] = \ 'projects/experimental/models/test_model' hook_instance = mock_hook.return_value hook_instance.create_job.side_effect = HttpError( resp=httplib2.Response({ 'status': http_error_code }), content=b'Forbidden') with self.assertRaises(HttpError) as context: prediction_task = MLEngineBatchPredictionOperator( job_id='test_prediction', project_id='test-project', region=input_with_model['region'], data_format=input_with_model['dataFormat'], input_paths=input_with_model['inputPaths'], output_path=input_with_model['outputPath'], model_name=input_with_model['modelName'].split('/')[-1], dag=self.dag, task_id='test-prediction') prediction_task.execute(None) mock_hook.assert_called_with('google_cloud_default', None) hook_instance.create_job.assert_called_with( 'test-project', { 'jobId': 'test_prediction', 'predictionInput': input_with_model }, ANY) self.assertEqual(http_error_code, context.exception.resp.status)
def testSuccessWithModel(self): with patch('airflow.contrib.operators.mlengine_operator.MLEngineHook') \ as mock_hook: input_with_model = self.INPUT_MISSING_ORIGIN.copy() input_with_model['modelName'] = \ 'projects/test-project/models/test_model' success_message = self.SUCCESS_MESSAGE_MISSING_INPUT.copy() success_message['predictionInput'] = input_with_model hook_instance = mock_hook.return_value hook_instance.get_job.side_effect = errors.HttpError( resp=httplib2.Response({ 'status': 404 }), content=b'some bytes') hook_instance.create_job.return_value = success_message prediction_task = MLEngineBatchPredictionOperator( job_id='test_prediction', project_id='test-project', region=input_with_model['region'], data_format=input_with_model['dataFormat'], input_paths=input_with_model['inputPaths'], output_path=input_with_model['outputPath'], model_name=input_with_model['modelName'].split('/')[-1], dag=self.dag, task_id='test-prediction') prediction_output = prediction_task.execute(None) mock_hook.assert_called_with('google_cloud_default', None) hook_instance.create_job.assert_called_once_with( 'test-project', { 'jobId': 'test_prediction', 'predictionInput': input_with_model }, ANY) self.assertEquals( success_message['predictionOutput'], prediction_output)