SSHHook expose auth_timeout parameter (apache#43048)

* SSHHook expose auth_timeout parameter * add tests
PaulKobow7536 · Oct 24, 2024 · 8e9ec41 · 8e9ec41
1 parent 7846ea8
commit 8e9ec41
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/providers/src/airflow/providers/ssh/hooks/ssh.py b/providers/src/airflow/providers/ssh/hooks/ssh.py
@@ -74,6 +74,7 @@ class SSHHook(BaseHook):
         iterable of algorithm identifiers, which will be disabled for the
         lifetime of the transport
     :param ciphers: list of ciphers to use in order of preference
+    :param auth_timeout: timeout (in seconds) for the attempt to authenticate with the remote_host
     """
 
     # List of classes to try loading private keys as, ordered (roughly) by most common to least common
@@ -121,6 +122,7 @@ def __init__(
         banner_timeout: float = 30.0,
         disabled_algorithms: dict | None = None,
         ciphers: list[str] | None = None,
+        auth_timeout: int | None = None,
     ) -> None:
         super().__init__()
         self.ssh_conn_id = ssh_conn_id
@@ -138,6 +140,7 @@ def __init__(
         self.disabled_algorithms = disabled_algorithms
         self.ciphers = ciphers
         self.host_proxy_cmd = None
+        self.auth_timeout = auth_timeout
 
         # Default values, overridable from Connection
         self.compress = True
@@ -332,6 +335,7 @@ def get_conn(self) -> paramiko.SSHClient:
             "sock": self.host_proxy,
             "look_for_keys": self.look_for_keys,
             "banner_timeout": self.banner_timeout,
+            "auth_timeout": self.auth_timeout,
         }
 
         if self.password:

diff --git a/providers/tests/ssh/hooks/test_ssh.py b/providers/tests/ssh/hooks/test_ssh.py
@@ -362,6 +362,7 @@ def test_ssh_connection_with_password(self, ssh_mock):
                 port="port",
                 sock=None,
                 look_for_keys=True,
+                auth_timeout=None,
             )
 
     @mock.patch("airflow.providers.ssh.hooks.ssh.paramiko.SSHClient")
@@ -381,6 +382,7 @@ def test_ssh_connection_without_password(self, ssh_mock):
                 port="port",
                 sock=None,
                 look_for_keys=True,
+                auth_timeout=None,
             )
 
     @mock.patch("airflow.providers.ssh.hooks.ssh.SSHTunnelForwarder")
@@ -569,6 +571,7 @@ def test_ssh_connection_with_private_key_extra(self, ssh_mock):
                 port="port",
                 sock=None,
                 look_for_keys=True,
+                auth_timeout=None,
             )
 
     @mock.patch("airflow.providers.ssh.hooks.ssh.paramiko.SSHClient")
@@ -592,6 +595,7 @@ def test_ssh_connection_with_private_key_passphrase_extra(self, ssh_mock):
                 port="port",
                 sock=None,
                 look_for_keys=True,
+                auth_timeout=None,
             )
 
     @mock.patch("airflow.providers.ssh.hooks.ssh.paramiko.SSHClient")
@@ -654,6 +658,7 @@ def test_ssh_connection_with_conn_timeout(self, ssh_mock):
             password="password",
             conn_timeout=20,
             key_file="fake.file",
+            auth_timeout=10,
         )
 
         with hook.get_conn():
@@ -668,6 +673,7 @@ def test_ssh_connection_with_conn_timeout(self, ssh_mock):
                 port="port",
                 sock=None,
                 look_for_keys=True,
+                auth_timeout=10,
             )
 
     @mock.patch("airflow.providers.ssh.hooks.ssh.paramiko.SSHClient")
@@ -695,6 +701,7 @@ def test_ssh_connection_with_conn_timeout_and_timeout(self, ssh_mock):
                 port="port",
                 sock=None,
                 look_for_keys=True,
+                auth_timeout=None,
             )
 
     @mock.patch("airflow.providers.ssh.hooks.ssh.paramiko.SSHClient")
@@ -718,6 +725,7 @@ def test_ssh_connection_with_timeout_extra(self, ssh_mock):
                 port="port",
                 sock=None,
                 look_for_keys=True,
+                auth_timeout=None,
             )
 
     @mock.patch("airflow.providers.ssh.hooks.ssh.paramiko.SSHClient")
@@ -743,6 +751,7 @@ def test_ssh_connection_with_conn_timeout_extra(self, ssh_mock):
                 port="port",
                 sock=None,
                 look_for_keys=True,
+                auth_timeout=None,
             )
 
     @mock.patch("airflow.providers.ssh.hooks.ssh.paramiko.SSHClient")
@@ -768,6 +777,7 @@ def test_ssh_connection_with_timeout_extra_and_conn_timeout_extra(self, ssh_mock
                 port="port",
                 sock=None,
                 look_for_keys=True,
+                auth_timeout=None,
             )
 
     @pytest.mark.parametrize(
@@ -837,6 +847,7 @@ def test_ssh_connection_with_all_timeout_param_and_extra_combinations(
                 port="port",
                 sock=None,
                 look_for_keys=True,
+                auth_timeout=None,
             )
 
     @pytest.mark.parametrize(
@@ -901,6 +912,7 @@ def test_ssh_with_extra_disabled_algorithms(self, ssh_mock):
                 sock=None,
                 look_for_keys=True,
                 disabled_algorithms=TEST_DISABLED_ALGORITHMS,
+                auth_timeout=None,
             )
 
     @mock.patch("airflow.providers.ssh.hooks.ssh.paramiko.SSHClient")