Python http_hook.HttpHook类代码示例

 def setUp(self):
     session = requests.Session()
     adapter = requests_mock.Adapter()
     session.mount('mock', adapter)
     self.get_hook = HttpHook(method='GET')
     self.post_hook = HttpHook(method='POST')
     configuration.load_test_config()

 def test_host_encoded_https_connection(self, mock_get_connection):
     c = Connection(conn_id='http_default', conn_type='http',
                    host='https://localhost')
     mock_get_connection.return_value = c
     hook = HttpHook()
     hook.get_conn({})
     self.assertEqual(hook.base_url, 'https://localhost')

 def execute(self, context):
     http = HttpHook(self.method, http_conn_id=self.http_conn_id)
     logging.info("Calling HTTP method")
     response = http.run(self.endpoint,
                         self.data,
                         self.headers,
                         self.extra_options)
     if self.response_check:
         if not self.response_check(response):
             raise AirflowException("Response check returned False.")

    def __init__(self,
                 endpoint,
                 http_conn_id='http_default',
                 method='GET',
                 request_params=None,
                 headers=None,
                 response_check=None,
                 extra_options=None, *args, **kwargs):
        super(HttpSensor, self).__init__(*args, **kwargs)
        self.endpoint = endpoint
        self.http_conn_id = http_conn_id
        self.request_params = request_params or {}
        self.headers = headers or {}
        self.extra_options = extra_options or {}
        self.response_check = response_check

        self.hook = HttpHook(
            method=method,
            http_conn_id=http_conn_id)

class TestHttpHook(unittest.TestCase):
    """Test get, post and raise_for_status"""
    def setUp(self):
        session = requests.Session()
        adapter = requests_mock.Adapter()
        session.mount('mock', adapter)
        self.get_hook = HttpHook(method='GET')
        self.post_hook = HttpHook(method='POST')
        configuration.load_test_config()

    @requests_mock.mock()
    def test_raise_for_status_with_200(self, m):

        m.get(
            'http://test:8080/v1/test',
            status_code=200,
            text='{"status":{"status": 200}}',
            reason='OK'
        )
        with mock.patch(
            'airflow.hooks.base_hook.BaseHook.get_connection',
            side_effect=get_airflow_connection
        ):

            resp = self.get_hook.run('v1/test')
            self.assertEqual(resp.text, '{"status":{"status": 200}}')

    @requests_mock.mock()
    @mock.patch('requests.Request')
    def test_get_request_with_port(self, m, request_mock):
        from requests.exceptions import MissingSchema

        with mock.patch(
            'airflow.hooks.base_hook.BaseHook.get_connection',
            side_effect=get_airflow_connection_with_port
        ):
            expected_url = 'http://test.com:1234/some/endpoint'
            for endpoint in ['some/endpoint', '/some/endpoint']:

                try:
                    self.get_hook.run(endpoint)
                except MissingSchema:
                    pass

                request_mock.assert_called_once_with(
                    mock.ANY,
                    expected_url,
                    headers=mock.ANY,
                    params=mock.ANY
                )

                request_mock.reset_mock()

    @requests_mock.mock()
    def test_get_request_do_not_raise_for_status_if_check_response_is_false(self, m):

        m.get(
            'http://test:8080/v1/test',
            status_code=404,
            text='{"status":{"status": 404}}',
            reason='Bad request'
        )

        with mock.patch(
            'airflow.hooks.base_hook.BaseHook.get_connection',
            side_effect=get_airflow_connection
        ):
            resp = self.get_hook.run('v1/test', extra_options={'check_response': False})
            self.assertEqual(resp.text, '{"status":{"status": 404}}')

    @requests_mock.mock()
    def test_hook_contains_header_from_extra_field(self, m):
        with mock.patch(
            'airflow.hooks.base_hook.BaseHook.get_connection',
            side_effect=get_airflow_connection
        ):
            expected_conn = get_airflow_connection()
            conn = self.get_hook.get_conn()
            self.assertDictContainsSubset(json.loads(expected_conn.extra), conn.headers)
            self.assertEqual(conn.headers.get('bareer'), 'test')

    @requests_mock.mock()
    def test_hook_uses_provided_header(self, m):
        conn = self.get_hook.get_conn(headers={"bareer": "newT0k3n"})
        self.assertEqual(conn.headers.get('bareer'), "newT0k3n")

    @requests_mock.mock()
    def test_hook_has_no_header_from_extra(self, m):
        conn = self.get_hook.get_conn()
        self.assertIsNone(conn.headers.get('bareer'))

    @requests_mock.mock()
    def test_hooks_header_from_extra_is_overridden(self, m):
        with mock.patch(
            'airflow.hooks.base_hook.BaseHook.get_connection',
            side_effect=get_airflow_connection
        ):
            conn = self.get_hook.get_conn(headers={"bareer": "newT0k3n"})
            self.assertEqual(conn.headers.get('bareer'), 'newT0k3n')

#.........这里部分代码省略.........

class HttpSensor(BaseSensorOperator):
    """
    Executes a HTTP get statement and returns False on failure:
        404 not found or response_check function returned False

    :param http_conn_id: The connection to run the sensor against
    :type http_conn_id: string
    :param method: The HTTP request method to use
    :type method: string
    :param endpoint: The relative part of the full url
    :type endpoint: string
    :param request_params: The parameters to be added to the GET url
    :type request_params: a dictionary of string key/value pairs
    :param headers: The HTTP headers to be added to the GET request
    :type headers: a dictionary of string key/value pairs
    :param response_check: A check against the 'requests' response object.
        Returns True for 'pass' and False otherwise.
    :type response_check: A lambda or defined function.
    :param extra_options: Extra options for the 'requests' library, see the
        'requests' documentation (options to modify timeout, ssl, etc.)
    :type extra_options: A dictionary of options, where key is string and value
        depends on the option that's being modified.
    """

    template_fields = ('endpoint', 'request_params')

    @apply_defaults
    def __init__(self,
                 endpoint,
                 http_conn_id='http_default',
                 method='GET',
                 request_params=None,
                 headers=None,
                 response_check=None,
                 extra_options=None, *args, **kwargs):
        super(HttpSensor, self).__init__(*args, **kwargs)
        self.endpoint = endpoint
        self.http_conn_id = http_conn_id
        self.request_params = request_params or {}
        self.headers = headers or {}
        self.extra_options = extra_options or {}
        self.response_check = response_check

        self.hook = HttpHook(
            method=method,
            http_conn_id=http_conn_id)

    def poke(self, context):
        self.log.info('Poking: %s', self.endpoint)
        try:
            response = self.hook.run(self.endpoint,
                                     data=self.request_params,
                                     headers=self.headers,
                                     extra_options=self.extra_options)
            if self.response_check:
                # run content check on response
                return self.response_check(response)
        except AirflowException as ae:
            if str(ae).startswith("404"):
                return False

            raise ae

        return True

def read_connection(**context):
    context['task_instance'].xcom_push(key='my_key', value='my_value')
    hook = HttpHook(http_conn_id='http_default')
    session = hook.get_conn(None)
    print("requests.sessions.Session: ", session)

def send_schedules_to_screen(**kwargs):  # hope it's op_kwargs + context
    data = json.dumps(kwargs['task_instance'].xcom_pull(task_ids='create_schedules'))
    http = HttpHook(kwargs['method'], http_conn_id=kwargs['http_conn_id'])
    return http.run(kwargs['endpoint'], data, kwargs['headers'])  # response_check from operator is on by default, I guess