def test_http_error(self):
        http_error_code = 403

        with patch('airflow.gcp.operators.mlengine.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_once_with('google_cloud_default', None)
                hook_instance.create_job.assert_called_once_with(
                    'test-project', {
                        'jobId': 'test_prediction',
                        'predictionInput': input_with_model
                    }, ANY)

            self.assertEqual(http_error_code, context.exception.resp.status)
    def test_success_with_model(self):
        with patch('airflow.gcp.operators.mlengine.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 = 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_once_with('google_cloud_default', None)
            hook_instance.create_job.assert_called_once_with(
                'test-project', {
                    'jobId': 'test_prediction',
                    'predictionInput': input_with_model
                }, ANY)
            self.assertEqual(success_message['predictionOutput'],
                             prediction_output)
    def test_success_with_uri(self, mock_hook):
        input_with_uri = self.INPUT_MISSING_ORIGIN.copy()
        input_with_uri['uri'] = 'gs://my_bucket/my_models/savedModel'
        success_message = self.SUCCESS_MESSAGE_MISSING_INPUT.copy()
        success_message['predictionInput'] = input_with_uri

        hook_instance = mock_hook.return_value
        hook_instance.get_job.side_effect = 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_uri['region'],
            data_format=input_with_uri['dataFormat'],
            input_paths=input_with_uri['inputPaths'],
            output_path=input_with_uri['outputPath'],
            uri=input_with_uri['uri'],
            dag=self.dag,
            task_id='test-prediction')
        prediction_output = prediction_task.execute(None)

        mock_hook.assert_called_once_with('google_cloud_default', None)
        hook_instance.create_job.assert_called_once_with(
            project_id='test-project',
            job={
                'jobId': 'test_prediction',
                'predictionInput': input_with_uri
            },
            use_existing_job_fn=ANY)
        self.assertEqual(success_message['predictionOutput'],
                         prediction_output)