[airflow] branch main updated: Add test_connection method for sftp hook (#19609)

[ur...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

uranusjr pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new ccb8095  Add test_connection method for sftp hook (#19609)
ccb8095 is described below

commit ccb809550db95089178385184cc898f5412cd266
Author: Constance Martineau <co...@gmail.com>
AuthorDate: Tue Nov 16 01:01:47 2021 -0500

    Add test_connection method for sftp hook (#19609)
---
 airflow/providers/sftp/hooks/sftp.py    |  9 +++++++++
 tests/providers/sftp/hooks/test_sftp.py | 31 +++++++++++++++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/airflow/providers/sftp/hooks/sftp.py b/airflow/providers/sftp/hooks/sftp.py
index 61b36cb..ec1f860 100644
--- a/airflow/providers/sftp/hooks/sftp.py
+++ b/airflow/providers/sftp/hooks/sftp.py
@@ -320,3 +320,12 @@ class SFTPHook(SSHHook):
         )
 
         return files, dirs, unknowns
+
+    def test_connection(self) -> Tuple[bool, str]:
+        """Test the SFTP connection by checking if remote entity '/some/path' exists"""
+        try:
+            conn = self.get_conn()
+            conn.pwd
+            return True, "Connection successfully tested"
+        except Exception as e:
+            return False, str(e)
diff --git a/tests/providers/sftp/hooks/test_sftp.py b/tests/providers/sftp/hooks/test_sftp.py
index 445d50e..35a5256 100644
--- a/tests/providers/sftp/hooks/test_sftp.py
+++ b/tests/providers/sftp/hooks/test_sftp.py
@@ -304,6 +304,37 @@ class TestSFTPHook(unittest.TestCase):
         assert dirs == [os.path.join(TMP_PATH, TMP_DIR_FOR_TESTS, SUB_DIR)]
         assert unknowns == []
 
+    @mock.patch('airflow.providers.sftp.hooks.sftp.SFTPHook.get_connection')
+    @mock.patch(
+        'airflow.providers.sftp.hooks.sftp.SFTPHook.get_conn.pwd', side_effect=Exception('Connection Error')
+    )
+    def test_connection_failure(self, mock_get_connection, mock_pwd):
+        connection = Connection(
+            login='login',
+            host='host',
+        )
+        mock_get_connection.return_value = connection
+
+        hook = SFTPHook()
+        status, msg = hook.test_connection()
+        assert status is False
+        assert msg == 'Connection Error'
+
+    @mock.patch('airflow.providers.sftp.hooks.sftp.SFTPHook.get_connection')
+    @mock.patch('airflow.providers.sftp.hooks.sftp.SFTPHook.get_conn.pwd')
+    def test_connection_success(self, mock_get_connection, mock_pwd):
+        connection = Connection(
+            login='login',
+            host='host',
+        )
+        mock_get_connection.return_value = connection
+        mock_pwd.return_value = '/home/some_user'
+
+        hook = SFTPHook()
+        status, msg = hook.test_connection()
+        assert status is True
+        assert msg == 'Connection successfully tested'
+
     def tearDown(self):
         shutil.rmtree(os.path.join(TMP_PATH, TMP_DIR_FOR_TESTS))
         os.remove(os.path.join(TMP_PATH, TMP_FILE_FOR_TESTS))