Fixup SSH connector tests for Python 3.12

pyinfra-dev · May 12, 2024 · 42a3e4e · 42a3e4e
1 parent 9d67e9b
commit 42a3e4e
Showing 1 changed file with 15 additions and 13 deletions.
diff --git a/tests/test_connectors/test_ssh.py b/tests/test_connectors/test_ssh.py
@@ -1054,14 +1054,17 @@ def test_get_sftp_fail(self, fake_sftp_client, fake_ssh_client):
                 )
 
     @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,
             gaierror,
             socket_error,
             EOFError,
         ):
+            fake_sleep.reset_mock()
+            fake_ssh_client.reset_mock()
+
             inventory = make_inventory(
                 hosts=("unresposivehost",), override_data={"ssh_connect_retries": 1}
             )
@@ -1070,24 +1073,26 @@ def test_ssh_connect_fail_retry(self, fake_sleep, fake_ssh_client):
             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,
             gaierror,
             socket_error,
             EOFError,
         ):
+            fake_sleep.reset_mock()
+            fake_ssh_client.reset_mock()
+
             inventory = make_inventory(
                 hosts=("unresposivehost",), override_data={"ssh_connect_retries": 1}
             )
@@ -1096,11 +1101,8 @@ def test_ssh_connect_fail_success(self, fake_sleep, fake_ssh_client):
             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