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 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 test_connection_without_host(self, mock_get_connection): c = Connection(conn_id='http_default', conn_type='http') mock_get_connection.return_value = c hook = HttpHook() hook.get_conn({}) self.assertEqual(hook.base_url, 'http://')
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')
def query_job_status(job_id, http_hook: HttpHook, token): session = http_hook.get_conn(headers={'Accept': "application/json"}) res = session.get(url=http_hook.base_url + "spider/api/job-status", params={ 'access_token': token, 'job_id': job_id }) return json.loads(res.content)
def prepare_connection(conn_id, route): http_hook = HttpHook(http_conn_id=conn_id) session = http_hook.get_conn() url = "/".join( [ u.strip("/") for u in [http_hook.base_url, session.headers["endpoint"], route] ] ) return http_hook, session, url
def post_status_info(context): CONN_ID = "process_report" ROUTE = "status" CRYPT_PRIVATE_KEY = "process_report_private_key" CRYPT_ALGORITHM = "process_report_crypt_algorithm" try: # Checking connection http_hook = HttpHook(http_conn_id=CONN_ID) session = http_hook.get_conn() endpoint = session.headers["endpoint"] url = http_hook.base_url.rstrip("/") + '/' + endpoint.strip("/") + "/" + ROUTE.lstrip("/") # Preparing data dag_run = context["dag_run"] data_format = "%Y-%m-%d %H:%M:%S" data = {"dag_id": dag_run.dag_id, "run_id": dag_run.run_id, "execution_date": dag_run.execution_date.strftime(data_format) if dag_run.execution_date else None, "start_date": dag_run.start_date.strftime(data_format) if dag_run.start_date else None, "end_date": dag_run.end_date.strftime(data_format) if dag_run.end_date else None, "state": dag_run.state, "tasks": []} for ti in dag_run.get_task_instances(): data["tasks"].append({"task_id": ti.task_id, "start_date": ti.start_date.strftime(data_format) if ti.start_date else None, "end_date": ti.end_date.strftime(data_format) if ti.end_date else None, "state": ti.state, "try_number": ti.try_number, "max_tries": ti.max_tries}) # Add DagRun results # Results should be collected only from CWLJobGatherer class or any other class that inherits from it. # Optimal solution would be to check ifinstance(ti, CWLJobGatherer), but unless code in cwlutils.py is # refactored, import of CWLJobGatherer may cause LOOP. # Therefore we check if isinstance(ti.xcom_pull(task_ids=ti.task_id), tuple), because only CWLJobGatherer # returns tuple. # Additionally, when post_status_info function is called as callback from CWLDAG, the context is equal to # the context of tasks[-1], where the order of tasks depends on DB and cannot be stable. if dag_run.state == State.SUCCESS: try: data["results"] = [ti.xcom_pull(task_ids=ti.task_id)[0] for ti in dag_run.get_task_instances() if isinstance(ti.xcom_pull(task_ids=ti.task_id), tuple)][0] except Exception as ex: print("Failed to collect results\n", ex) # Try to encode data if CRYPT_PRIVATE_KEY is set in Variable try: data = jwt.encode(data, Variable.get(CRYPT_PRIVATE_KEY), algorithm=Variable.get(CRYPT_ALGORITHM)).decode("utf-8") except Exception as e: print("Failed to encrypt status data:\n", e) # Posting results prepped_request = session.prepare_request(requests.Request("POST", url, json={"payload": data})) http_hook.run_and_check(session, prepped_request, {}) except Exception as e: print("Failed to POST status updates:\n", e)
def upload(self, model, desc): http = HttpHook('POST', http_conn_id=self.http_conn_id) session = http.get_conn(self.headers) req = Request('POST', urljoin(http.base_url, self.endpoint), data={'description': desc}, files={'file': model}, headers=self.headers) prepped_request = session.prepare_request(req) response = http.run_and_check(session, prepped_request, self.extra_options) return response.text
def _get_jwt_token(self): http_auth = HttpHook('POST', http_conn_id=self.http_conn_id) tv = http_auth.get_conn().auth tk = ('username', 'password') payload_auth = dict(zip(tk, tv)) self.log.info("Calling HTTP auth method ") response = http_auth.run(self.endpoint_auth, json.dumps(payload_auth), self.headers, self.extra_options) if self.response_check: if not self.response_check(response): raise AirflowException("Response check returned False.") resp_json = response.json() return resp_json['access_token']
def add_tasks(tasks, task_hook_conn_id): task_hook = HttpHook(http_conn_id=task_hook_conn_id) data = {'tasks': tasks} logging.info("增加[%s]个任务", len(tasks)) conn = task_hook.get_connection(task_hook.http_conn_id) session = task_hook.get_conn(headers={}) res = session.post(url=task_hook.base_url + "task/addbatch", json=data, headers={ 'Content-Type': 'application/json', 'spider-client-id': conn.login, 'spider-api-key': conn.password }) res = json.loads(res.content) succeed = isinstance(res, list) and len(res) > 0 if not succeed: raise Exception(res)
def execute(self, context): http = HttpHook(self.method, http_conn_id=self.http_conn_id) session = http.get_conn(self.headers) with open(self.__evaluation_path, 'rb') as evaluation_file: evaluation_data = evaluation_file.read() req = Request(self.method, urljoin(http.base_url, self.endpoint), files={'file': evaluation_data}, headers=self.headers) response = http.run_and_check(session, session.prepare_request(req), self.extra_options) result = json.loads(response.text) if result['status']: self.log.info(f'Success. evaluation_id: {result["evaluation_id"]}') else: raise AirflowException(result['message']) return result["evaluation_id"]
def prepare_connection(conn_id, route=""): http_hook = HttpHook(http_conn_id=conn_id) session = http_hook.get_conn() url = http_hook.base_url.rstrip("/") + '/' + session.headers[ "endpoint"].strip("/") + "/" + route.lstrip("/") return http_hook, session, url
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') @requests_mock.mock() def test_post_request(self, m): m.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, m): m.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, m): m.post( 'http://*****:*****@mock.patch('airflow.hooks.http_hook.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=requests.exceptions.ConnectionError ) def send_and_raise(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 ) def test_header_from_extra_and_run_method_are_merged(self): def run_and_return(session, prepped_request, extra_options, **kwargs): return prepped_request # The job failed for some reason with mock.patch( 'airflow.hooks.http_hook.HttpHook.run_and_check', side_effect=run_and_return ): with mock.patch( 'airflow.hooks.base_hook.BaseHook.get_connection', side_effect=get_airflow_connection ): pr = self.get_hook.run('v1/test', headers={'some_other_header': 'test'}) actual = dict(pr.headers) self.assertEqual(actual.get('bareer'), 'test') self.assertEqual(actual.get('some_other_header'), 'test') @mock.patch('airflow.hooks.http_hook.HttpHook.get_connection') def test_http_connection(self, mock_get_connection): c = Connection(conn_id='http_default', conn_type='http', host='localhost', schema='http') mock_get_connection.return_value = c hook = HttpHook() hook.get_conn({}) self.assertEqual(hook.base_url, 'http://localhost') @mock.patch('airflow.hooks.http_hook.HttpHook.get_connection') def test_https_connection(self, mock_get_connection): c = Connection(conn_id='http_default', conn_type='http', host='localhost', schema='https') mock_get_connection.return_value = c hook = HttpHook() hook.get_conn({}) self.assertEqual(hook.base_url, 'https://localhost') @mock.patch('airflow.hooks.http_hook.HttpHook.get_connection') def test_host_encoded_http_connection(self, mock_get_connection): c = Connection(conn_id='http_default', conn_type='http', host='http://localhost') mock_get_connection.return_value = c hook = HttpHook() hook.get_conn({}) self.assertEqual(hook.base_url, 'http://localhost') @mock.patch('airflow.hooks.http_hook.HttpHook.get_connection') 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')
class SAPODataOperator(BaseOperator): """ SAPODataOperator for downloading data from SAP OData services and convert it to pandas DataFrame. Pushes result to xcom. :param http_hook: predefined http_hook to use for data transmitting. Either `http_hook` or `http_conn_id` needs to be provided. :type http_hook: airflow.hooks.http_hook.HttpHook :param http_conn_id: connection id from airflow Connections. `http_conn_id` will be ignored if `http_hook` is provided. :type http_conn_id: str :param service_url: URL of the OData service. It will replace the `base_url` which was defined in `http_hook` or predefined in the connection of `http_conn_id`. :type service_url: str :param function: name of the function, that provides the data. :type function: str Either `function` or `entity` needs to be provided. :param entity: name of the entity set that provides the data. :type entity: str :param property: property for data recieving. Nullable. If not provided, will be replaced with `entity` value. :type property: str :param parameters: parameters for the service. :type parameters: dict """ # template_fields = ('parameters') ui_color = '#dae2f2' ui_fgcolor = '#19477e' @apply_defaults def __init__(self, http_hook=None, http_conn_id=None, service_url=None, function=None, entity=None, property=None, parameters=dict(), *args, **kwargs): super().__init__(*args, **kwargs) self.http_hook = http_hook self.http_conn_id = http_conn_id self.service_url = service_url self.function = function self.entity = entity self.property = property self.parameters = parameters if not self.property and not self.function: raise AirflowException( f"Either `function` or `entity` needs to be provided.") if not self.service_url: raise AirflowException(f"`service_url` needs to be provided.") if not isinstance(self.parameters, dict): raise TypeError( f"`parameters` should be dictionary like {parameter: value}, got {type(self.parameters)}" ) def execute(self, context): step_msg = None result = None try: if self.http_conn_id: if self.http_hook and isinstance(self.http_hook, HttpHook): self.log.info( "`http_conn_id` is ignored when `http_hook` is provided." ) else: self.log.info("`http_hook` is not provided or invalid. " "Trying `http_conn_id` to create HttpHook.") self.http_hook = HttpHook(http_conn_id=self.http_conn_id) if not self.http_hook: step_msg = "creating http session" raise AirflowException( "Cannot operate without `http_hook` or `http_conn_id`.") else: url = requests.urllib3.util.url.parse_url(self.service_url) self.http_hook.base_url = f"{url.scheme}://{url.host}" + f"{':' + str(url.port) if url.port else ''}" with self.http_hook.get_conn() as sess: client = pyodata.Client(self.service_url, sess) if self.function is not None and isinstance( self.function, str): if self.entity is not None: self.log.info( "entity is ignored when function is provided.") step_msg = f"preparing data request for function {self.function}" request = client.functions.__getattr__(self.function) if self.parameters: for name, value in self.parameters.items(): request = request.parameter(name, value) step_msg = f"sending request and recieving data" request = request.execute() step_msg = f"creating result pandas DataFrame" result = pd.DataFrame.from_dict(request['results']) elif self.entity: if not self.property: self.log.info( "`property` is not provided, filling with `entity` value" ) self.property = self.entity step_msg = f"preparing data request for entity set {self.entity}" data = client.entity_sets.__getattr__(self.entity)\ .get_entity(**self.parameters).nav(self.property).get_entities() step_msg = f"sending request and recieving data" data = data.execute() step_msg = f"recieving headers" columns = { es.name: es for es in client.schema.entity_sets }[self.property].entity_type.proprties() step_msg = f"creating result pandas DataFrame" result = pd.DataFrame.from_records([{ col.name: d.__getattr__(col.name) for col in columns } for d in data]) except Exception as e: raise AirflowException( f"Error while step \"{step_msg}\", error: {str(e)}") return result
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)
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://*****:*****@mock.patch('requests.Session') @mock.patch('requests.Request') def test_get_request_with_port(self, request_mock, session_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() session_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() @mock.patch('requests.Request') def test_hook_with_method_in_lowercase(self, m, request_mock): from requests.exceptions import MissingSchema, InvalidURL 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, 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') @requests_mock.mock() def test_post_request(self, m): m.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, m): m.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, m): m.post('http://*****:*****@mock.patch('airflow.hooks.http_hook.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(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(u'http://*****:*****@mock.patch('airflow.hooks.http_hook.HttpHook.get_connection') def test_http_connection(self, mock_get_connection): c = Connection(conn_id='http_default', conn_type='http', host='localhost', schema='http') mock_get_connection.return_value = c hook = HttpHook() hook.get_conn({}) self.assertEqual(hook.base_url, 'http://localhost') @mock.patch('airflow.hooks.http_hook.HttpHook.get_connection') def test_https_connection(self, mock_get_connection): c = Connection(conn_id='http_default', conn_type='http', host='localhost', schema='https') mock_get_connection.return_value = c hook = HttpHook() hook.get_conn({}) self.assertEqual(hook.base_url, 'https://localhost') @mock.patch('airflow.hooks.http_hook.HttpHook.get_connection') def test_host_encoded_http_connection(self, mock_get_connection): c = Connection(conn_id='http_default', conn_type='http', host='http://localhost') mock_get_connection.return_value = c hook = HttpHook() hook.get_conn({}) self.assertEqual(hook.base_url, 'http://localhost') @mock.patch('airflow.hooks.http_hook.HttpHook.get_connection') 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 test_method_converted_to_uppercase_when_created_in_lowercase(self): self.assertEqual(self.get_lowercase_hook.method, 'GET') @mock.patch('airflow.hooks.http_hook.HttpHook.get_connection') def test_connection_without_host(self, mock_get_connection): c = Connection(conn_id='http_default', conn_type='http') mock_get_connection.return_value = c 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)
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.assertEquals(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.assertEquals(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.assertEquals(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.assertEquals(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.assertEquals(conn.headers.get('bareer'), 'newT0k3n') @requests_mock.mock() def test_post_request(self, m): m.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.assertEquals(resp.status_code, 200) @requests_mock.mock() def test_post_request_with_error_code(self, m): m.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, m): m.post( 'http://*****:*****@mock.patch('airflow.hooks.http_hook.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=requests.exceptions.ConnectionError ) def send_and_raise(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.assertEquals( self.get_hook._retry_obj.stop.max_attempt_number + 1, mocked_session.call_count ) def test_header_from_extra_and_run_method_are_merged(self): def run_and_return(session, prepped_request, extra_options, **kwargs): return prepped_request # The job failed for some reason with mock.patch( 'airflow.hooks.http_hook.HttpHook.run_and_check', side_effect=run_and_return ): with mock.patch( 'airflow.hooks.base_hook.BaseHook.get_connection', side_effect=get_airflow_connection ): pr = self.get_hook.run('v1/test', headers={'some_other_header': 'test'}) actual = dict(pr.headers) self.assertEquals(actual.get('bareer'), 'test') self.assertEquals(actual.get('some_other_header'), 'test') @mock.patch('airflow.hooks.http_hook.HttpHook.get_connection') def test_http_connection(self, mock_get_connection): c = models.Connection(conn_id='http_default', conn_type='http', host='localhost', schema='http') mock_get_connection.return_value = c hook = HttpHook() hook.get_conn({}) self.assertEqual(hook.base_url, 'http://localhost') @mock.patch('airflow.hooks.http_hook.HttpHook.get_connection') def test_https_connection(self, mock_get_connection): c = models.Connection(conn_id='http_default', conn_type='http', host='localhost', schema='https') mock_get_connection.return_value = c hook = HttpHook() hook.get_conn({}) self.assertEqual(hook.base_url, 'https://localhost') @mock.patch('airflow.hooks.http_hook.HttpHook.get_connection') def test_host_encoded_http_connection(self, mock_get_connection): c = models.Connection(conn_id='http_default', conn_type='http', host='http://localhost') mock_get_connection.return_value = c hook = HttpHook() hook.get_conn({}) self.assertEqual(hook.base_url, 'http://localhost') @mock.patch('airflow.hooks.http_hook.HttpHook.get_connection') def test_host_encoded_https_connection(self, mock_get_connection): c = models.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 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)