Skip to content

Commit

Permalink
Add test connection method to http hook (#16568)
Browse files Browse the repository at this point in the history
  • Loading branch information
msumit committed Jun 22, 2021
1 parent 19ed074 commit 61fdf84
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 1 deletion.
10 changes: 9 additions & 1 deletion airflow/providers/http/hooks/http.py
Expand Up @@ -95,7 +95,7 @@ def get_conn(self, headers: Optional[Dict[Any, Any]] = None) -> requests.Session

def run(
self,
endpoint: Optional[str],
endpoint: Optional[str] = None,
data: Optional[Union[Dict[str, Any], str]] = None,
headers: Optional[Dict[str, Any]] = None,
extra_options: Optional[Dict[str, Any]] = None,
Expand Down Expand Up @@ -227,3 +227,11 @@ def run_with_advanced_retry(self, _retry_args: Dict[Any, Any], *args: Any, **kwa
self._retry_obj = tenacity.Retrying(**_retry_args)

return self._retry_obj(self.run, *args, **kwargs)

def test_connection(self):
"""Test HTTP Connection"""
try:
self.run()
return True, 'Connection successfully tested'
except Exception as e: # noqa pylint: disable=broad-except
return False, str(e)
16 changes: 16 additions & 0 deletions tests/providers/http/hooks/test_http.py
Expand Up @@ -353,5 +353,21 @@ def test_verify_false_parameter_overwrites_set_requests_ca_bundle_env_var(self,
verify=False,
)

@requests_mock.mock()
def test_connection_success(self, m):
m.get('http://test:8080', status_code=200, json={"status": {"status": 200}}, reason='OK')
with mock.patch('airflow.hooks.base.BaseHook.get_connection', side_effect=get_airflow_connection):
status, msg = self.get_hook.test_connection()
assert status is True
assert msg == 'Connection successfully tested'

@requests_mock.mock()
def test_connection_failure(self, m):
m.get('http://test:8080', status_code=500, json={"message": "internal server error"}, reason='NOT_OK')
with mock.patch('airflow.hooks.base.BaseHook.get_connection', side_effect=get_airflow_connection):
status, msg = self.get_hook.test_connection()
assert status is False
assert msg == '500:NOT_OK'


send_email_test = mock.Mock()

0 comments on commit 61fdf84

Please sign in to comment.