Improve exception raised when failing to open SFTP sessions.

Author: Fizzadar

pyinfra/api/connectors/ssh.py

@@ -9,6 +9,7 @@
 )
 
 import click
+import six
 
 from paramiko import (
     AuthenticationException,
@@ -306,7 +307,14 @@ def run_shell_command(
 @memoize
 def _get_sftp_connection(host):
     transport = host.connection.get_transport()
-    return SFTPClient.from_transport(transport)
+
+    try:
+        return SFTPClient.from_transport(transport)
+    except SSHException as e:
+        six.raise_from(ConnectError((
+            'Unable to establish SFTP connection. Check that the SFTP subsystem '
+            'for the SSH service at {0} is enabled.'
+        ).format(host)), e)
 
 
 def _get_file(host, remote_filename, filename_or_io):

tests/test_connectors/test_ssh.py

@@ -20,7 +20,7 @@
 from pyinfra.api import Config, MaskString, State, StringCommand
 from pyinfra.api.connect import connect_all
 from pyinfra.api.connectors.ssh import _get_sftp_connection
-from pyinfra.api.exceptions import PyinfraError
+from pyinfra.api.exceptions import ConnectError, PyinfraError
 
 from ..util import make_inventory
 
@@ -650,3 +650,24 @@ def test_get_file_su_user(self, fake_sftp_client, fake_ssh_client):
         fake_sftp_client.from_transport().getfo.assert_called_with(
             '/tmp/pyinfra-e9c0d3c8ffca943daa0e75511b0a09c84b59c508', fake_open(),
         )
+
+    @patch('pyinfra.api.connectors.ssh.SSHClient')
+    @patch('pyinfra.api.connectors.ssh.SFTPClient')
+    def test_get_sftp_fail(self, fake_sftp_client, fake_ssh_client):
+        inventory = make_inventory(hosts=('anotherhost',))
+        State(inventory, Config())
+        host = inventory.get_host('anotherhost')
+        host.connect()
+
+        def raise_exception(*args, **kwargs):
+            raise SSHException()
+
+        fake_sftp_client.from_transport.side_effect = raise_exception
+
+        fake_open = mock_open(read_data='test!')
+        with patch('pyinfra.api.util.open', fake_open, create=True):
+            with self.assertRaises(ConnectError):
+                host.put_file(
+                    'not-a-file', 'not-another-file',
+                    print_output=True,
+                )