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)
