This is an automated email from the ASF dual-hosted git repository.
kaxilnaik 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 6751c9635f8 Improve speed of SSH & SFTP tests (#45938)
6751c9635f8 is described below
commit 6751c9635f86838086bbdfc52a17ad6818fbd03a
Author: Kaxil Naik <kaxiln...@apache.org>
AuthorDate: Wed Jan 22 23:52:27 2025 +0530
Improve speed of SSH & SFTP tests (#45938)
Before:
> 4 passed, 1 warning in 64.03s (0:01:04)
After:
> 4 passed, 1 warning in 4.60s
To run them:
```
pytest
providers/tests/sftp/operators/test_sftp.py::TestSFTPOperator::test_arg_checking
providers/tests/ssh/hooks/test_ssh.py::TestSSHHook::test_command_timeout_not_set
providers/tests/ssh/hooks/test_ssh.py::TestSSHHook::test_command_timeout_success
providers/tests/ssh/hooks/test_ssh.py::TestSSHHook::test_command_timeout_fail
```
---
providers/src/airflow/providers/ssh/hooks/ssh.py | 9 ++++-----
providers/tests/sftp/operators/test_sftp.py | 6 +++++-
providers/tests/ssh/hooks/test_ssh.py | 19 ++++++++++---------
3 files changed, 19 insertions(+), 15 deletions(-)
diff --git a/providers/src/airflow/providers/ssh/hooks/ssh.py
b/providers/src/airflow/providers/ssh/hooks/ssh.py
index 3502459c644..c88f2a1de8e 100644
--- a/providers/src/airflow/providers/ssh/hooks/ssh.py
+++ b/providers/src/airflow/providers/ssh/hooks/ssh.py
@@ -37,7 +37,6 @@ from airflow.hooks.base import BaseHook
from airflow.utils.platform import getuser
from airflow.utils.types import NOTSET, ArgNotSet
-TIMEOUT_DEFAULT = 10
CMD_TIMEOUT = 10
@@ -113,7 +112,7 @@ class SSHHook(BaseHook):
key_file: str | None = None,
port: int | None = None,
conn_timeout: int | None = None,
- cmd_timeout: int | ArgNotSet | None = NOTSET,
+ cmd_timeout: float | ArgNotSet | None = NOTSET,
keepalive_interval: int = 30,
banner_timeout: float = 30.0,
disabled_algorithms: dict | None = None,
@@ -175,7 +174,7 @@ class SSHHook(BaseHook):
if "cmd_timeout" in extra_options and self.cmd_timeout is
NOTSET:
if extra_options["cmd_timeout"]:
- self.cmd_timeout = int(extra_options["cmd_timeout"])
+ self.cmd_timeout = float(extra_options["cmd_timeout"])
else:
self.cmd_timeout = None
@@ -427,11 +426,11 @@ class SSHHook(BaseHook):
command: str,
get_pty: bool,
environment: dict | None,
- timeout: int | ArgNotSet | None = NOTSET,
+ timeout: float | ArgNotSet | None = NOTSET,
) -> tuple[int, bytes, bytes]:
self.log.info("Running command: %s", command)
- cmd_timeout: int | None
+ cmd_timeout: float | None
if not isinstance(timeout, ArgNotSet):
cmd_timeout = timeout
elif not isinstance(self.cmd_timeout, ArgNotSet):
diff --git a/providers/tests/sftp/operators/test_sftp.py
b/providers/tests/sftp/operators/test_sftp.py
index 34bd5b9b6ed..c0d2c03db8e 100644
--- a/providers/tests/sftp/operators/test_sftp.py
+++ b/providers/tests/sftp/operators/test_sftp.py
@@ -314,7 +314,11 @@ class TestSFTPOperator:
assert content_received == self.test_remote_file_content
@mock.patch.dict("os.environ", {"AIRFLOW_CONN_" + TEST_CONN_ID.upper():
"ssh://test_id@localhost"})
- def test_arg_checking(self):
+
@mock.patch("airflow.providers.sftp.operators.sftp.SFTPHook.retrieve_directory")
+ @mock.patch("airflow.providers.sftp.operators.sftp.SFTPHook.retrieve_file")
+
@mock.patch("airflow.providers.sftp.operators.sftp.SFTPHook.create_directory")
+ @mock.patch("airflow.providers.sftp.operators.sftp.SFTPHook.store_file")
+ def test_arg_checking(self, mock_store_file, mock_create_dir,
mock_retrieve_file, mock_retrieve_dir):
dag = DAG(
dag_id="unit_tests_sftp_op_arg_checking",
schedule=None,
diff --git a/providers/tests/ssh/hooks/test_ssh.py
b/providers/tests/ssh/hooks/test_ssh.py
index e09f2eeee0a..d778a46411b 100644
--- a/providers/tests/ssh/hooks/test_ssh.py
+++ b/providers/tests/ssh/hooks/test_ssh.py
@@ -829,30 +829,28 @@ class TestSSHHook:
)
assert ret == (0, b"airflow\n", b"")
- @pytest.mark.flaky(reruns=5)
def test_command_timeout_success(self):
hook = SSHHook(
ssh_conn_id="ssh_default",
conn_timeout=30,
- cmd_timeout=15,
+ cmd_timeout=2,
banner_timeout=100,
)
with hook.get_conn() as client:
ret = hook.exec_ssh_client_command(
client,
- "sleep 10; echo airflow",
+ "sleep 0.1; echo airflow",
False,
None,
)
assert ret == (0, b"airflow\n", b"")
- @pytest.mark.flaky(reruns=5)
def test_command_timeout_fail(self):
hook = SSHHook(
ssh_conn_id="ssh_default",
conn_timeout=30,
- cmd_timeout=5,
+ cmd_timeout=0.001,
banner_timeout=100,
)
@@ -860,12 +858,12 @@ class TestSSHHook:
with pytest.raises(AirflowException):
hook.exec_ssh_client_command(
client,
- "sleep 10",
+ "sleep 1",
False,
None,
)
- def test_command_timeout_not_set(self):
+ def test_command_timeout_not_set(self, monkeypatch):
hook = SSHHook(
ssh_conn_id="ssh_default",
conn_timeout=30,
@@ -873,12 +871,15 @@ class TestSSHHook:
banner_timeout=100,
)
+ # Mock the timeout to not wait for that many seconds
+ monkeypatch.setattr("airflow.providers.ssh.hooks.ssh.CMD_TIMEOUT",
0.001)
+
with hook.get_conn() as client:
- # sleeping for 20 sec which is longer than default timeout of 10
seconds
+ # sleeping for more secs than default timeout
# to validate that no timeout is applied
hook.exec_ssh_client_command(
client,
- "sleep 20",
+ "sleep 0.1",
environment=False,
get_pty=None,
)
