diff options
| -rw-r--r-- | paramiko/sftp_client.py | 3 | ||||
| -rwxr-xr-x | tests/test_sftp.py | 21 |
2 files changed, 23 insertions, 1 deletions
diff --git a/paramiko/sftp_client.py b/paramiko/sftp_client.py index 3e85a8c9..b45ad442 100644 --- a/paramiko/sftp_client.py +++ b/paramiko/sftp_client.py @@ -39,6 +39,7 @@ from paramiko.sftp import BaseSFTP, CMD_OPENDIR, CMD_HANDLE, SFTPError, CMD_READ from paramiko.sftp_attr import SFTPAttributes from paramiko.ssh_exception import SSHException from paramiko.sftp_file import SFTPFile +from paramiko.util import ClosingContextManager def _to_unicode(s): @@ -58,7 +59,7 @@ def _to_unicode(s): b_slash = b'/' -class SFTPClient(BaseSFTP): +class SFTPClient(BaseSFTP, ClosingContextManager): """ SFTP client object. diff --git a/tests/test_sftp.py b/tests/test_sftp.py index 1ae9781d..e3f7e59b 100755 --- a/tests/test_sftp.py +++ b/tests/test_sftp.py @@ -195,6 +195,18 @@ class SFTPTest (unittest.TestCase): pass sftp = paramiko.SFTP.from_transport(tc) + def test_2_sftp_can_be_used_as_context_manager(self): + """ + verify that the sftp session is closed when exiting the context manager + """ + global sftp + with sftp: + pass + try: + self._assert_opening_file_raises_error(sftp) + finally: + sftp = paramiko.SFTP.from_transport(tc) + def test_3_write(self): """ verify that a file can be created and written, and the size is correct. @@ -796,6 +808,15 @@ class SFTPTest (unittest.TestCase): sftp.remove('%s/nonutf8data' % FOLDER) + def _assert_opening_file_raises_error(self, sftp): + try: + sftp.open(FOLDER + '/test2', 'w') + self.fail('expected exception') + # TODO: extract failing command (the open) to share with a passing assertion + except EOFError as error: + pass + + if __name__ == '__main__': SFTPTest.init_loopback() # logging is required by test_N_file_with_percent |