test_mlengine_operator.py 文件源码

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)