Skip to content

Commit

Permalink
Fix pytests
Browse files Browse the repository at this point in the history
  • Loading branch information
@jscheffl
jscheffl committed Dec 2, 2024
1 parent 20c2145 commit cfe73bd
Showing 1 changed file with 3 additions and 176 deletions.
179 changes: 3 additions & 176 deletions providers/tests/ssh/hooks/test_ssh.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@
import pytest

from airflow import settings
from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
from airflow.exceptions import AirflowException
from airflow.models import Connection
from airflow.providers.ssh.hooks.ssh import SSHHook
from airflow.utils import db
Expand Down Expand Up @@ -676,180 +676,6 @@ def test_ssh_connection_with_conn_timeout(self, ssh_mock):
auth_timeout=10,
)

@mock.patch("airflow.providers.ssh.hooks.ssh.paramiko.SSHClient")
def test_ssh_connection_with_conn_timeout_and_timeout(self, ssh_mock):
with pytest.warns(AirflowProviderDeprecationWarning, match=".*Please use `conn_timeout` instead..*"):
hook = SSHHook(
remote_host="remote_host",
port="port",
username="username",
password="password",
timeout=10,
conn_timeout=20,
key_file="fake.file",
)

with hook.get_conn():
ssh_mock.return_value.connect.assert_called_once_with(
banner_timeout=30.0,
hostname="remote_host",
username="username",
password="password",
key_filename="fake.file",
timeout=20,
compress=True,
port="port",
sock=None,
look_for_keys=True,
auth_timeout=None,
)

@mock.patch("airflow.providers.ssh.hooks.ssh.paramiko.SSHClient")
def test_ssh_connection_with_timeout_extra(self, ssh_mock):
with pytest.warns(AirflowProviderDeprecationWarning, match=".*Please use `conn_timeout` instead..*"):
hook = SSHHook(
ssh_conn_id=self.CONN_SSH_WITH_TIMEOUT_EXTRA,
remote_host="remote_host",
port="port",
username="username",
timeout=10,
)

with hook.get_conn():
ssh_mock.return_value.connect.assert_called_once_with(
banner_timeout=30.0,
hostname="remote_host",
username="username",
timeout=20,
compress=True,
port="port",
sock=None,
look_for_keys=True,
auth_timeout=None,
)

@mock.patch("airflow.providers.ssh.hooks.ssh.paramiko.SSHClient")
def test_ssh_connection_with_conn_timeout_extra(self, ssh_mock):
with pytest.warns(AirflowProviderDeprecationWarning, match=".*Please use `conn_timeout` instead..*"):
hook = SSHHook(
ssh_conn_id=self.CONN_SSH_WITH_CONN_TIMEOUT_EXTRA,
remote_host="remote_host",
port="port",
username="username",
timeout=10,
conn_timeout=15,
)

# conn_timeout parameter wins over extra options
with hook.get_conn():
ssh_mock.return_value.connect.assert_called_once_with(
banner_timeout=30.0,
hostname="remote_host",
username="username",
timeout=15,
compress=True,
port="port",
sock=None,
look_for_keys=True,
auth_timeout=None,
)

@mock.patch("airflow.providers.ssh.hooks.ssh.paramiko.SSHClient")
def test_ssh_connection_with_timeout_extra_and_conn_timeout_extra(self, ssh_mock):
with pytest.warns(AirflowProviderDeprecationWarning, match=".*Please use `conn_timeout` instead..*"):
hook = SSHHook(
ssh_conn_id=self.CONN_SSH_WITH_TIMEOUT_AND_CONN_TIMEOUT_EXTRA,
remote_host="remote_host",
port="port",
username="username",
timeout=10,
conn_timeout=15,
)

# conn_timeout parameter wins over extra options
with hook.get_conn():
ssh_mock.return_value.connect.assert_called_once_with(
banner_timeout=30.0,
hostname="remote_host",
username="username",
timeout=15,
compress=True,
port="port",
sock=None,
look_for_keys=True,
auth_timeout=None,
)

@pytest.mark.parametrize(
"timeout, conn_timeout, timeoutextra, conn_timeoutextra, expected_value",
[
(TEST_TIMEOUT, TEST_CONN_TIMEOUT, True, True, TEST_CONN_TIMEOUT),
(TEST_TIMEOUT, TEST_CONN_TIMEOUT, True, False, TEST_CONN_TIMEOUT),
(TEST_TIMEOUT, TEST_CONN_TIMEOUT, False, True, TEST_CONN_TIMEOUT),
(TEST_TIMEOUT, TEST_CONN_TIMEOUT, False, False, TEST_CONN_TIMEOUT),
(TEST_TIMEOUT, None, True, True, TEST_CONN_TIMEOUT),
(TEST_TIMEOUT, None, True, False, TEST_TIMEOUT),
(TEST_TIMEOUT, None, False, True, TEST_CONN_TIMEOUT),
(TEST_TIMEOUT, None, False, False, TEST_TIMEOUT),
(None, TEST_CONN_TIMEOUT, True, True, TEST_CONN_TIMEOUT),
(None, TEST_CONN_TIMEOUT, True, False, TEST_CONN_TIMEOUT),
(None, TEST_CONN_TIMEOUT, False, True, TEST_CONN_TIMEOUT),
(None, TEST_CONN_TIMEOUT, False, False, TEST_CONN_TIMEOUT),
(None, None, True, True, TEST_CONN_TIMEOUT),
(None, None, True, False, TEST_TIMEOUT),
(None, None, False, True, TEST_CONN_TIMEOUT),
(None, None, False, False, 10),
],
)
@mock.patch("airflow.providers.ssh.hooks.ssh.paramiko.SSHClient")
def test_ssh_connection_with_all_timeout_param_and_extra_combinations(
self, ssh_mock, timeout, conn_timeout, timeoutextra, conn_timeoutextra, expected_value
):
if timeoutextra and conn_timeoutextra:
ssh_conn_id = self.CONN_SSH_WITH_TIMEOUT_AND_CONN_TIMEOUT_EXTRA
elif timeoutextra and not conn_timeoutextra:
ssh_conn_id = self.CONN_SSH_WITH_TIMEOUT_EXTRA
elif not timeoutextra and conn_timeoutextra:
ssh_conn_id = self.CONN_SSH_WITH_CONN_TIMEOUT_EXTRA
else:
ssh_conn_id = self.CONN_SSH_WITH_NO_EXTRA

if timeout or timeoutextra:
with pytest.warns(
AirflowProviderDeprecationWarning, match=".*Please use `conn_timeout` instead..*"
):
hook = SSHHook(
ssh_conn_id=ssh_conn_id,
remote_host="remote_host",
port="port",
username="username",
timeout=timeout,
conn_timeout=conn_timeout,
)
else:
hook = SSHHook(
ssh_conn_id=ssh_conn_id,
remote_host="remote_host",
port="port",
username="username",
timeout=timeout,
conn_timeout=conn_timeout,
)

# conn_timeout parameter wins over extra options
with hook.get_conn():
ssh_mock.return_value.connect.assert_called_once_with(
banner_timeout=30.0,
hostname="remote_host",
username="username",
timeout=expected_value,
compress=True,
port="port",
sock=None,
look_for_keys=True,
auth_timeout=None,
)

@pytest.mark.parametrize(
"cmd_timeout, cmd_timeoutextra, null_cmd_timeoutextra, expected_value",
[
Expand Down Expand Up @@ -899,6 +725,7 @@ def test_ssh_with_extra_disabled_algorithms(self, ssh_mock):
remote_host="remote_host",
port="port",
username="username",
conn_timeout=42,
)

with hook.get_conn():
Expand All @@ -907,7 +734,7 @@ def test_ssh_with_extra_disabled_algorithms(self, ssh_mock):
hostname="remote_host",
username="username",
compress=True,
timeout=10,
timeout=42,
port="port",
sock=None,
look_for_keys=True,
Expand Down

0 comments on commit cfe73bd

Please sign in to comment.