예제 #1
0
    def test_connection_without_host(self, mock_get_connection):
        conn = Connection(conn_id='http_default', conn_type='http')
        mock_get_connection.return_value = conn

        hook = HttpHook()
        hook.get_conn({})
        self.assertEqual(hook.base_url, 'http://')
예제 #2
0
 def test_host_encoded_https_connection(self, mock_get_connection):
     conn = Connection(conn_id='http_default',
                       conn_type='http',
                       host='https://localhost')
     mock_get_connection.return_value = conn
     hook = HttpHook()
     hook.get_conn({})
     self.assertEqual(hook.base_url, 'https://localhost')
예제 #3
0
 def test_https_connection(self, mock_get_connection):
     conn = Connection(conn_id='http_default',
                       conn_type='http',
                       host='localhost',
                       schema='https')
     mock_get_connection.return_value = conn
     hook = HttpHook()
     hook.get_conn({})
     assert hook.base_url == 'https://localhost'
예제 #4
0
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.get_lowercase_hook = HttpHook(method='get')
        self.post_hook = HttpHook(method='POST')

    @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.Session')
    @mock.patch('requests.Request')
    def test_get_request_with_port(self, mock_requests, request_mock,
                                   mock_session):
        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, mock_requests):
        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()
    @mock.patch('requests.Request')
    def test_hook_with_method_in_lowercase(self, mock_requests, request_mock):
        from requests.exceptions import InvalidURL, MissingSchema

        with mock.patch('airflow.hooks.base_hook.BaseHook.get_connection',
                        side_effect=get_airflow_connection_with_port):
            data = "test params"
            try:
                self.get_lowercase_hook.run('v1/test', data=data)
            except (MissingSchema, InvalidURL):
                pass
            request_mock.assert_called_once_with(mock.ANY,
                                                 mock.ANY,
                                                 headers=mock.ANY,
                                                 params=data)

    @requests_mock.mock()
    def test_hook_uses_provided_header(self, mock_requests):
        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, mock_requests):
        conn = self.get_hook.get_conn()
        self.assertIsNone(conn.headers.get('bareer'))

    @requests_mock.mock()
    def test_hooks_header_from_extra_is_overridden(self, mock_requests):
        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')

    @requests_mock.mock()
    def test_post_request(self, mock_requests):
        mock_requests.post('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.post_hook.run('v1/test')
            self.assertEqual(resp.status_code, 200)

    @requests_mock.mock()
    def test_post_request_with_error_code(self, mock_requests):
        mock_requests.post(
            'http://test:8080/v1/test',
            status_code=418,
            text='{"status":{"status": 418}}',
            reason='I\'m a teapot',
        )

        with mock.patch('airflow.hooks.base_hook.BaseHook.get_connection',
                        side_effect=get_airflow_connection):
            with self.assertRaises(AirflowException):
                self.post_hook.run('v1/test')

    @requests_mock.mock()
    def test_post_request_do_not_raise_for_status_if_check_response_is_false(
            self, mock_requests):
        mock_requests.post(
            'http://*****:*****@mock.patch('airflow.providers.http.hooks.http.requests.Session')
    def test_retry_on_conn_error(self, mocked_session):

        retry_args = dict(
            wait=tenacity.wait_none(),
            stop=tenacity.stop_after_attempt(7),
            retry=tenacity.retry_if_exception_type(
                requests.exceptions.ConnectionError),
        )

        def send_and_raise(unused_request, **kwargs):
            raise requests.exceptions.ConnectionError

        mocked_session().send.side_effect = send_and_raise
        # The job failed for some reason
        with self.assertRaises(tenacity.RetryError):
            self.get_hook.run_with_advanced_retry(endpoint='v1/test',
                                                  _retry_args=retry_args)
        self.assertEqual(self.get_hook._retry_obj.stop.max_attempt_number + 1,
                         mocked_session.call_count)

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

        m.get('http://*****:*****@mock.patch('airflow.providers.http.hooks.http.HttpHook.get_connection')
    def test_http_connection(self, mock_get_connection):
        conn = Connection(conn_id='http_default',
                          conn_type='http',
                          host='localhost',
                          schema='http')
        mock_get_connection.return_value = conn
        hook = HttpHook()
        hook.get_conn({})
        self.assertEqual(hook.base_url, 'http://localhost')

    @mock.patch('airflow.providers.http.hooks.http.HttpHook.get_connection')
    def test_https_connection(self, mock_get_connection):
        conn = Connection(conn_id='http_default',
                          conn_type='http',
                          host='localhost',
                          schema='https')
        mock_get_connection.return_value = conn
        hook = HttpHook()
        hook.get_conn({})
        self.assertEqual(hook.base_url, 'https://localhost')

    @mock.patch('airflow.providers.http.hooks.http.HttpHook.get_connection')
    def test_host_encoded_http_connection(self, mock_get_connection):
        conn = Connection(conn_id='http_default',
                          conn_type='http',
                          host='http://localhost')
        mock_get_connection.return_value = conn
        hook = HttpHook()
        hook.get_conn({})
        self.assertEqual(hook.base_url, 'http://localhost')

    @mock.patch('airflow.providers.http.hooks.http.HttpHook.get_connection')
    def test_host_encoded_https_connection(self, mock_get_connection):
        conn = Connection(conn_id='http_default',
                          conn_type='http',
                          host='https://localhost')
        mock_get_connection.return_value = conn
        hook = HttpHook()
        hook.get_conn({})
        self.assertEqual(hook.base_url, 'https://localhost')

    def test_method_converted_to_uppercase_when_created_in_lowercase(self):
        self.assertEqual(self.get_lowercase_hook.method, 'GET')

    @mock.patch('airflow.providers.http.hooks.http.HttpHook.get_connection')
    def test_connection_without_host(self, mock_get_connection):
        conn = Connection(conn_id='http_default', conn_type='http')
        mock_get_connection.return_value = conn

        hook = HttpHook()
        hook.get_conn({})
        self.assertEqual(hook.base_url, 'http://')

    @parameterized.expand([
        'GET',
        'POST',
    ])
    @requests_mock.mock()
    def test_json_request(self, method, mock_requests):
        obj1 = {'a': 1, 'b': 'abc', 'c': [1, 2, {"d": 10}]}

        def match_obj1(request):
            return request.json() == obj1

        mock_requests.request(method=method,
                              url='//test:8080/v1/test',
                              additional_matcher=match_obj1)

        with mock.patch('airflow.hooks.base_hook.BaseHook.get_connection',
                        side_effect=get_airflow_connection):
            # will raise NoMockAddress exception if obj1 != request.json()
            HttpHook(method=method).run('v1/test', json=obj1)