pyinfra 2.9.1__py2.py3-none-any.whl → 3.0__py2.py3-none-any.whl

tests/test_connectors/test_local.py

@@ -45,18 +45,17 @@ class TestLocalConnector(TestCase):
         command = "echo Šablony"
         self.fake_popen_mock().returncode = 0
 
-        out = host.run_shell_command(command, stdin="hello", print_output=True)
-        assert len(out) == 3
+        out = host.run_shell_command(command, _stdin="hello", print_output=True)
+        assert len(out) == 2
 
-        status, stdout, stderr = out
+        status, output = out
         assert status is True
         self.fake_popen_mock().stdin.write.assert_called_with(b"hello\n")
 
         combined_out = host.run_shell_command(
             command,
-            stdin="hello",
+            _stdin="hello",
             print_output=True,
-            return_combined_output=True,
         )
         assert len(combined_out) == 2
 
@@ -79,9 +78,9 @@ class TestLocalConnector(TestCase):
         self.fake_popen_mock().returncode = 0
 
         out = host.run_shell_command(command, print_output=True, print_input=True)
-        assert len(out) == 3
+        assert len(out) == 2
 
-        status, stdout, stderr = out
+        status, output = out
         assert status is True
 
         self.fake_popen_mock.assert_called_with(
@@ -105,8 +104,8 @@ class TestLocalConnector(TestCase):
         command = "echo hi"
         self.fake_popen_mock().returncode = 1
 
-        out = host.run_shell_command(command, success_exit_codes=[1])
-        assert len(out) == 3
+        out = host.run_shell_command(command, _success_exit_codes=[1])
+        assert len(out) == 2
         assert out[0] is True
 
     def test_run_shell_command_error(self):
@@ -118,7 +117,7 @@ class TestLocalConnector(TestCase):
         self.fake_popen_mock().returncode = 1
 
         out = host.run_shell_command(command)
-        assert len(out) == 3
+        assert len(out) == 2
         assert out[0] is False
 
     def test_put_file(self):
@@ -210,7 +209,7 @@ class TestLocalConnector(TestCase):
         command = "echo Šablony"
         self.fake_popen_mock().returncode = 0
 
-        host.run_shell_command(command, stdin=["hello", "abc"], print_output=True)
+        host.run_shell_command(command, _stdin=["hello", "abc"], print_output=True)
         self.fake_popen_mock().stdin.write.assert_has_calls(
             [
                 call(b"hello\n"),
@@ -226,7 +225,7 @@ class TestLocalConnector(TestCase):
         command = "echo Šablony"
         self.fake_popen_mock().returncode = 0
 
-        host.run_shell_command(command, stdin=StringIO("hello\nabc"), print_output=True)
+        host.run_shell_command(command, _stdin=StringIO("hello\nabc"), print_output=True)
         self.fake_popen_mock().stdin.write.assert_has_calls(
             [
                 call(b"hello\n"),

tests/test_connectors/test_ssh.py

@@ -10,7 +10,6 @@ import pyinfra
 from pyinfra.api import Config, MaskString, State, StringCommand
 from pyinfra.api.connect import connect_all
 from pyinfra.api.exceptions import ConnectError, PyinfraError
-from pyinfra.connectors.ssh import _get_sftp_connection
 
 from ..util import make_inventory
 
@@ -28,7 +27,6 @@ class TestSSHConnector(TestCase):
     def setUp(self):
         self.fake_connect_patch = patch("pyinfra.connectors.ssh.SSHClient.connect")
         self.fake_connect_mock = self.fake_connect_patch.start()
-        _get_sftp_connection.cache = {}
 
     def tearDown(self):
         self.fake_connect_patch.stop()
@@ -58,8 +56,8 @@ class TestSSHConnector(TestCase):
 
         assert len(state.active_hosts) == 2
 
-    @patch("pyinfra.connectors.ssh.path.isfile", lambda *args, **kwargs: True)
-    @patch("pyinfra.connectors.ssh.RSAKey.from_private_key_file")
+    @patch("pyinfra.connectors.ssh_util.path.isfile", lambda *args, **kwargs: True)
+    @patch("pyinfra.connectors.ssh_util.RSAKey.from_private_key_file")
     def test_connect_exceptions(self, fake_key_open):
         for exception_class in (
             AuthenticationException,
@@ -83,8 +81,8 @@ class TestSSHConnector(TestCase):
     def test_connect_with_rsa_ssh_key(self):
         state = State(make_inventory(hosts=(("somehost", {"ssh_key": "testkey"}),)), Config())
 
-        with patch("pyinfra.connectors.ssh.path.isfile", lambda *args, **kwargs: True), patch(
-            "pyinfra.connectors.ssh.RSAKey.from_private_key_file",
+        with patch("pyinfra.connectors.ssh_util.path.isfile", lambda *args, **kwargs: True), patch(
+            "pyinfra.connectors.ssh_util.RSAKey.from_private_key_file",
         ) as fake_key_open:
             fake_key = MagicMock()
             fake_key_open.return_value = fake_key
@@ -104,10 +102,10 @@ class TestSSHConnector(TestCase):
                 pkey=fake_key,
                 timeout=10,
                 username="vagrant",
-                _pyinfra_ssh_forward_agent=None,
+                _pyinfra_ssh_forward_agent=False,
                 _pyinfra_ssh_config_file=None,
                 _pyinfra_ssh_known_hosts_file=None,
-                _pyinfra_ssh_strict_host_key_checking=None,
+                _pyinfra_ssh_strict_host_key_checking="accept-new",
                 _pyinfra_ssh_paramiko_connect_kwargs=None,
             )
 
@@ -128,8 +126,8 @@ class TestSSHConnector(TestCase):
             Config(),
         )
 
-        with patch("pyinfra.connectors.ssh.path.isfile", lambda *args, **kwargs: True), patch(
-            "pyinfra.connectors.ssh.RSAKey.from_private_key_file",
+        with patch("pyinfra.connectors.ssh_util.path.isfile", lambda *args, **kwargs: True), patch(
+            "pyinfra.connectors.ssh_util.RSAKey.from_private_key_file",
         ) as fake_key_open:
             fake_key = MagicMock()
 
@@ -150,11 +148,11 @@ class TestSSHConnector(TestCase):
     def test_connect_with_rsa_ssh_key_password_from_prompt(self):
         state = State(make_inventory(hosts=(("somehost", {"ssh_key": "testkey"}),)), Config())
 
-        with patch("pyinfra.connectors.ssh.path.isfile", lambda *args, **kwargs: True), patch(
-            "pyinfra.connectors.ssh.getpass",
+        with patch("pyinfra.connectors.ssh_util.path.isfile", lambda *args, **kwargs: True), patch(
+            "pyinfra.connectors.ssh_util.getpass",
             lambda *args, **kwargs: "testpass",
         ), patch(
-            "pyinfra.connectors.ssh.RSAKey.from_private_key_file",
+            "pyinfra.connectors.ssh_util.RSAKey.from_private_key_file",
         ) as fake_key_open:
             fake_key = MagicMock()
 
@@ -177,8 +175,8 @@ class TestSSHConnector(TestCase):
     def test_connect_with_rsa_ssh_key_missing_password(self):
         state = State(make_inventory(hosts=(("somehost", {"ssh_key": "testkey"}),)), Config())
 
-        with patch("pyinfra.connectors.ssh.path.isfile", lambda *args, **kwargs: True), patch(
-            "pyinfra.connectors.ssh.RSAKey.from_private_key_file",
+        with patch("pyinfra.connectors.ssh_util.path.isfile", lambda *args, **kwargs: True), patch(
+            "pyinfra.connectors.ssh_util.RSAKey.from_private_key_file",
         ) as fake_key_open:
             fake_key_open.side_effect = make_raise_exception_function(PasswordRequiredException)
 
@@ -203,17 +201,17 @@ class TestSSHConnector(TestCase):
         fake_fail_from_private_key_file = MagicMock()
         fake_fail_from_private_key_file.side_effect = make_raise_exception_function(SSHException)
 
-        with patch("pyinfra.connectors.ssh.path.isfile", lambda *args, **kwargs: True), patch(
-            "pyinfra.connectors.ssh.DSSKey.from_private_key_file",
+        with patch("pyinfra.connectors.ssh_util.path.isfile", lambda *args, **kwargs: True), patch(
+            "pyinfra.connectors.ssh_util.DSSKey.from_private_key_file",
             fake_fail_from_private_key_file,
         ), patch(
-            "pyinfra.connectors.ssh.ECDSAKey.from_private_key_file",
+            "pyinfra.connectors.ssh_util.ECDSAKey.from_private_key_file",
             fake_fail_from_private_key_file,
         ), patch(
-            "pyinfra.connectors.ssh.Ed25519Key.from_private_key_file",
+            "pyinfra.connectors.ssh_util.Ed25519Key.from_private_key_file",
             fake_fail_from_private_key_file,
         ), patch(
-            "pyinfra.connectors.ssh.RSAKey.from_private_key_file",
+            "pyinfra.connectors.ssh_util.RSAKey.from_private_key_file",
         ) as fake_key_open:
 
             def fake_key_open_fail(*args, **kwargs):
@@ -236,10 +234,10 @@ class TestSSHConnector(TestCase):
     def test_connect_with_dss_ssh_key(self):
         state = State(make_inventory(hosts=(("somehost", {"ssh_key": "testkey"}),)), Config())
 
-        with patch("pyinfra.connectors.ssh.path.isfile", lambda *args, **kwargs: True), patch(
-            "pyinfra.connectors.ssh.RSAKey.from_private_key_file",
+        with patch("pyinfra.connectors.ssh_util.path.isfile", lambda *args, **kwargs: True), patch(
+            "pyinfra.connectors.ssh_util.RSAKey.from_private_key_file",
         ) as fake_rsa_key_open, patch(
-            "pyinfra.connectors.ssh.DSSKey.from_private_key_file",
+            "pyinfra.connectors.ssh_util.DSSKey.from_private_key_file",
         ) as fake_key_open:  # noqa
             fake_rsa_key_open.side_effect = make_raise_exception_function(SSHException)
 
@@ -259,10 +257,10 @@ class TestSSHConnector(TestCase):
                 pkey=fake_key,
                 timeout=10,
                 username="vagrant",
-                _pyinfra_ssh_forward_agent=None,
+                _pyinfra_ssh_forward_agent=False,
                 _pyinfra_ssh_config_file=None,
                 _pyinfra_ssh_known_hosts_file=None,
-                _pyinfra_ssh_strict_host_key_checking=None,
+                _pyinfra_ssh_strict_host_key_checking="accept-new",
                 _pyinfra_ssh_paramiko_connect_kwargs=None,
             )
 
@@ -283,10 +281,10 @@ class TestSSHConnector(TestCase):
             Config(),
         )
 
-        with patch("pyinfra.connectors.ssh.path.isfile", lambda *args, **kwargs: True), patch(
-            "pyinfra.connectors.ssh.RSAKey.from_private_key_file",
+        with patch("pyinfra.connectors.ssh_util.path.isfile", lambda *args, **kwargs: True), patch(
+            "pyinfra.connectors.ssh_util.RSAKey.from_private_key_file",
         ) as fake_rsa_key_open, patch(
-            "pyinfra.connectors.ssh.DSSKey.from_private_key_file",
+            "pyinfra.connectors.ssh_util.DSSKey.from_private_key_file",
         ) as fake_dss_key_open:  # noqa
 
             def fake_rsa_key_open_fail(*args, **kwargs):
@@ -318,10 +316,10 @@ class TestSSHConnector(TestCase):
                 pkey=fake_dss_key,
                 timeout=10,
                 username="vagrant",
-                _pyinfra_ssh_forward_agent=None,
+                _pyinfra_ssh_forward_agent=False,
                 _pyinfra_ssh_config_file=None,
                 _pyinfra_ssh_known_hosts_file=None,
-                _pyinfra_ssh_strict_host_key_checking=None,
+                _pyinfra_ssh_strict_host_key_checking="accept-new",
                 _pyinfra_ssh_paramiko_connect_kwargs=None,
             )
 
@@ -362,18 +360,17 @@ class TestSSHConnector(TestCase):
         command = "echo Šablony"
         fake_stdout.channel.recv_exit_status.return_value = 0
 
-        out = host.run_shell_command(command, stdin="hello", print_output=True)
-        assert len(out) == 3
+        out = host.run_shell_command(command, _stdin="hello", print_output=True)
+        assert len(out) == 2
 
-        status, stdout, stderr = out
+        status, output = out
         assert status is True
         fake_stdin.write.assert_called_with(b"hello\n")
 
         combined_out = host.run_shell_command(
             command,
-            stdin="hello",
+            _stdin="hello",
             print_output=True,
-            return_combined_output=True,
         )
         assert len(combined_out) == 2
 
@@ -397,9 +394,9 @@ class TestSSHConnector(TestCase):
         fake_stdout.channel.recv_exit_status.return_value = 0
 
         out = host.run_shell_command(command, print_output=True, print_input=True)
-        assert len(out) == 3
+        assert len(out) == 2
 
-        status, stdout, stderr = out
+        status, output = out
         assert status is True
 
         fake_ssh.exec_command.assert_called_with(
@@ -428,8 +425,8 @@ class TestSSHConnector(TestCase):
         command = "echo hi"
         fake_stdout.channel.recv_exit_status.return_value = 1
 
-        out = host.run_shell_command(command, success_exit_codes=[1])
-        assert len(out) == 3
+        out = host.run_shell_command(command, _success_exit_codes=[1])
+        assert len(out) == 2
         assert out[0] is True
 
     @patch("pyinfra.connectors.ssh.SSHClient")
@@ -449,19 +446,29 @@ class TestSSHConnector(TestCase):
         fake_stdout.channel.recv_exit_status.return_value = 1
 
         out = host.run_shell_command(command)
-        assert len(out) == 3
+        assert len(out) == 2
         assert out[0] is False
 
     @patch("pyinfra.connectors.util.getpass")
     @patch("pyinfra.connectors.ssh.SSHClient")
-    def test_run_shell_command_sudo_password_prompt(
+    def test_run_shell_command_sudo_password_automatic_prompt(
         self,
         fake_ssh_client,
         fake_getpass,
     ):
         fake_ssh = MagicMock()
-        fake_stdout = MagicMock()
-        fake_ssh.exec_command.return_value = MagicMock(), fake_stdout, MagicMock()
+        first_fake_stdout = MagicMock()
+        second_fake_stdout = MagicMock()
+        third_fake_stdout = MagicMock()
+
+        first_fake_stdout.__iter__.return_value = ["sudo: a password is required\r"]
+        second_fake_stdout.__iter__.return_value = ["/tmp/pyinfra-sudo-askpass-XXXXXXXXXXXX"]
+
+        fake_ssh.exec_command.side_effect = [
+            (MagicMock(), first_fake_stdout, MagicMock()),  # command w/o sudo password
+            (MagicMock(), second_fake_stdout, MagicMock()),  # SUDO_ASKPASS_COMMAND
+            (MagicMock(), third_fake_stdout, MagicMock()),  # command with sudo pw
+        ]
 
         fake_ssh_client.return_value = fake_ssh
         fake_getpass.return_value = "password"
@@ -472,15 +479,18 @@ class TestSSHConnector(TestCase):
         host.connect()
 
         command = "echo Šablony"
-        fake_stdout.__iter__.return_value = ["/tmp/pyinfra-sudo-askpass-XXXXXXXXXXXX"]
-        fake_stdout.channel.recv_exit_status.return_value = 0
+        first_fake_stdout.channel.recv_exit_status.return_value = 1
+        second_fake_stdout.channel.recv_exit_status.return_value = 0
+        third_fake_stdout.channel.recv_exit_status.return_value = 0
 
-        out = host.run_shell_command(command, sudo=True, use_sudo_password=True, print_output=True)
-        assert len(out) == 3
+        out = host.run_shell_command(command, _sudo=True, print_output=True)
+        assert len(out) == 2
 
-        status, stdout, stderr = out
+        status, output = out
         assert status is True
 
+        fake_ssh.exec_command.assert_any_call(("sudo -H -n sh -c 'echo Šablony'"), get_pty=False)
+
         fake_ssh.exec_command.assert_called_with(
             (
                 "env SUDO_ASKPASS=/tmp/pyinfra-sudo-askpass-XXXXXXXXXXXX "
@@ -492,7 +502,7 @@ class TestSSHConnector(TestCase):
 
     @patch("pyinfra.connectors.util.getpass")
     @patch("pyinfra.connectors.ssh.SSHClient")
-    def test_run_shell_command_sudo_password_automatic_prompt(
+    def test_run_shell_command_sudo_password_automatic_prompt_with_special_chars_in_password(
         self,
         fake_ssh_client,
         fake_getpass,
@@ -512,7 +522,7 @@ class TestSSHConnector(TestCase):
         ]
 
         fake_ssh_client.return_value = fake_ssh
-        fake_getpass.return_value = "password"
+        fake_getpass.return_value = "p@ss'word';"
 
         inventory = make_inventory(hosts=("somehost",))
         State(inventory, Config())
@@ -524,10 +534,10 @@ class TestSSHConnector(TestCase):
         second_fake_stdout.channel.recv_exit_status.return_value = 0
         third_fake_stdout.channel.recv_exit_status.return_value = 0
 
-        out = host.run_shell_command(command, sudo=True, print_output=True)
-        assert len(out) == 3
+        out = host.run_shell_command(command, _sudo=True, print_output=True)
+        assert len(out) == 2
 
-        status, stdout, stderr = out
+        status, output = out
         assert status is True
 
         fake_ssh.exec_command.assert_any_call(("sudo -H -n sh -c 'echo Šablony'"), get_pty=False)
@@ -535,7 +545,7 @@ class TestSSHConnector(TestCase):
         fake_ssh.exec_command.assert_called_with(
             (
                 "env SUDO_ASKPASS=/tmp/pyinfra-sudo-askpass-XXXXXXXXXXXX "
-                "PYINFRA_SUDO_PASSWORD=password "
+                """PYINFRA_SUDO_PASSWORD='p@ss'"'"'word'"'"';' """
                 "sudo -H -A -k sh -c 'echo Šablony'"
             ),
             get_pty=False,
@@ -545,13 +555,13 @@ class TestSSHConnector(TestCase):
     #
 
     @patch("pyinfra.connectors.ssh.SSHClient")
-    @patch("pyinfra.connectors.util._get_sudo_password")
+    @patch("pyinfra.connectors.util.getpass")
     def test_run_shell_command_retry_for_sudo_password(
         self,
-        fake_get_sudo_password,
+        fake_getpass,
         fake_ssh_client,
     ):
-        fake_get_sudo_password.return_value = "PASSWORD"
+        fake_getpass.return_value = "PASSWORD"
 
         fake_ssh = MagicMock()
         fake_stdin = MagicMock()
@@ -571,13 +581,13 @@ class TestSSHConnector(TestCase):
         return_values = [1, 0]  # return 0 on the second call
         fake_stdout.channel.recv_exit_status.side_effect = lambda: return_values.pop(0)
 
-        out = host.run_shell_command(command)
-        assert len(out) == 3
+        out = host.run_shell_command(command, _sudo=True)
+        assert len(out) == 2
         assert out[0] is True
-        assert fake_get_sudo_password.called
+        assert fake_getpass.called
         fake_ssh.exec_command.assert_called_with(
             "env SUDO_ASKPASS=/tmp/pyinfra-sudo-askpass-XXXXXXXXXXXX "
-            "PYINFRA_SUDO_PASSWORD=PASSWORD sh -c 'echo hi'",
+            "PYINFRA_SUDO_PASSWORD=PASSWORD sudo -H -A -k sh -c 'echo hi'",
             get_pty=False,
         )
 
@@ -624,8 +634,8 @@ class TestSSHConnector(TestCase):
                 "not-a-file",
                 "not another file",
                 print_output=True,
-                sudo=True,
-                sudo_user="ubuntu",
+                _sudo=True,
+                _sudo_user="ubuntu",
             )
 
         assert status is True
@@ -675,8 +685,8 @@ class TestSSHConnector(TestCase):
                 "not-a-file",
                 "not another file",
                 print_output=True,
-                doas=True,
-                doas_user="ubuntu",
+                _doas=True,
+                _doas_user="ubuntu",
             )
 
         assert status is True
@@ -726,7 +736,7 @@ class TestSSHConnector(TestCase):
                 "not-a-file",
                 "not-another-file",
                 print_output=True,
-                su_user="centos",
+                _su_user="centos",
             )
 
         assert status is False
@@ -763,7 +773,7 @@ class TestSSHConnector(TestCase):
                 "not-a-file",
                 "not-another-file",
                 print_output=True,
-                su_user="centos",
+                _su_user="centos",
             )
 
         assert status is False
@@ -808,8 +818,8 @@ class TestSSHConnector(TestCase):
                 "not-a-file",
                 "not another file",
                 print_output=True,
-                sudo=True,
-                sudo_user="ubuntu",
+                _sudo=True,
+                _sudo_user="ubuntu",
                 remote_temp_filename="/a-different-tempfile",
             )
 
@@ -865,8 +875,8 @@ class TestSSHConnector(TestCase):
                 "not-a-file",
                 "not-another-file",
                 print_output=True,
-                sudo=True,
-                sudo_user="ubuntu",
+                _sudo=True,
+                _sudo_user="ubuntu",
             )
 
         assert status is True
@@ -876,7 +886,7 @@ class TestSSHConnector(TestCase):
                 call(
                     (
                         "sudo -H -n -u ubuntu sh -c 'cp not-a-file "
-                        "/tmp/pyinfra-e9c0d3c8ffca943daa0e75511b0a09c84b59c508 && chmod +r not-a-file'"  # noqa
+                        "/tmp/pyinfra-e9c0d3c8ffca943daa0e75511b0a09c84b59c508 && chmod +r /tmp/pyinfra-e9c0d3c8ffca943daa0e75511b0a09c84b59c508'"  # noqa
                     ),
                     get_pty=False,
                 ),
@@ -910,8 +920,8 @@ class TestSSHConnector(TestCase):
             "not-a-file",
             "not-another-file",
             print_output=True,
-            sudo=True,
-            sudo_user="ubuntu",
+            _sudo=True,
+            _sudo_user="ubuntu",
         )
 
         assert status is False
@@ -921,7 +931,7 @@ class TestSSHConnector(TestCase):
                 call(
                     (
                         "sudo -H -n -u ubuntu sh -c 'cp not-a-file "
-                        "/tmp/pyinfra-e9c0d3c8ffca943daa0e75511b0a09c84b59c508 && chmod +r not-a-file'"  # noqa
+                        "/tmp/pyinfra-e9c0d3c8ffca943daa0e75511b0a09c84b59c508 && chmod +r /tmp/pyinfra-e9c0d3c8ffca943daa0e75511b0a09c84b59c508'"  # noqa
                     ),
                     get_pty=False,
                 ),
@@ -946,8 +956,8 @@ class TestSSHConnector(TestCase):
                 "not-a-file",
                 "not-another-file",
                 print_output=True,
-                sudo=True,
-                sudo_user="ubuntu",
+                _sudo=True,
+                _sudo_user="ubuntu",
             )
 
         assert status is False
@@ -957,7 +967,7 @@ class TestSSHConnector(TestCase):
                 call(
                     (
                         "sudo -H -n -u ubuntu sh -c 'cp not-a-file "
-                        "/tmp/pyinfra-e9c0d3c8ffca943daa0e75511b0a09c84b59c508 && chmod +r not-a-file'"  # noqa
+                        "/tmp/pyinfra-e9c0d3c8ffca943daa0e75511b0a09c84b59c508 && chmod +r /tmp/pyinfra-e9c0d3c8ffca943daa0e75511b0a09c84b59c508'"  # noqa
                     ),
                     get_pty=False,
                 ),
@@ -994,7 +1004,7 @@ class TestSSHConnector(TestCase):
                 "not-a-file",
                 "not-another-file",
                 print_output=True,
-                su_user="centos",
+                _su_user="centos",
             )
 
         assert status is True
@@ -1005,7 +1015,7 @@ class TestSSHConnector(TestCase):
                     (
                         "su centos -c 'sh -c '\"'\"'cp not-a-file "
                         "/tmp/pyinfra-e9c0d3c8ffca943daa0e75511b0a09c84b59c508 && chmod +r "
-                        "not-a-file'\"'\"''"
+                        "/tmp/pyinfra-e9c0d3c8ffca943daa0e75511b0a09c84b59c508'\"'\"''"
                     ),
                     get_pty=False,
                 ),
@@ -1044,7 +1054,7 @@ class TestSSHConnector(TestCase):
                 )
 
     @patch("pyinfra.connectors.ssh.SSHClient")
-    @patch("time.sleep")
+    @patch("pyinfra.connectors.ssh.sleep")
     def test_ssh_connect_fail_retry(self, fake_sleep, fake_ssh_client):
         for exception_class in (
             SSHException,
@@ -1052,6 +1062,9 @@ class TestSSHConnector(TestCase):
             socket_error,
             EOFError,
         ):
+            fake_sleep.reset_mock()
+            fake_ssh_client.reset_mock()
+
             inventory = make_inventory(
                 hosts=("unresposivehost",), override_data={"ssh_connect_retries": 1}
             )
@@ -1060,17 +1073,16 @@ class TestSSHConnector(TestCase):
             unresposivehost = inventory.get_host("unresposivehost")
             assert unresposivehost.data.ssh_connect_retries == 1
 
-            fake_ssh = MagicMock()
-            fake_ssh.connect.side_effect = exception_class()
-            fake_ssh_client.return_value = fake_ssh
+            fake_ssh_client().connect.side_effect = exception_class()
 
             with self.assertRaises(ConnectError):
                 unresposivehost.connect(show_errors=False, raise_exceptions=True)
-            assert fake_sleep.called_once()
-            assert fake_ssh_client.connect.called_twice()
+
+            fake_sleep.assert_called_once()
+            assert fake_ssh_client().connect.call_count == 2
 
     @patch("pyinfra.connectors.ssh.SSHClient")
-    @patch("time.sleep")
+    @patch("pyinfra.connectors.ssh.sleep")
     def test_ssh_connect_fail_success(self, fake_sleep, fake_ssh_client):
         for exception_class in (
             SSHException,
@@ -1078,6 +1090,9 @@ class TestSSHConnector(TestCase):
             socket_error,
             EOFError,
         ):
+            fake_sleep.reset_mock()
+            fake_ssh_client.reset_mock()
+
             inventory = make_inventory(
                 hosts=("unresposivehost",), override_data={"ssh_connect_retries": 1}
             )
@@ -1086,11 +1101,8 @@ class TestSSHConnector(TestCase):
             unresposivehost = inventory.get_host("unresposivehost")
             assert unresposivehost.data.ssh_connect_retries == 1
 
-            connection = MagicMock()
-            fake_ssh = MagicMock()
-            fake_ssh.connect.side_effect = [exception_class(), connection]
-            fake_ssh_client.return_value = fake_ssh
+            fake_ssh_client().connect.side_effect = [exception_class(), MagicMock()]
 
             unresposivehost.connect(show_errors=False, raise_exceptions=True)
-            assert fake_sleep.called_once()
-            assert fake_ssh_client.connect.called_twice()
+            fake_sleep.assert_called_once()
+            assert fake_ssh_client().connect.call_count == 2