mylib/scripts/sftp_test.py

@@ -1,13 +1,10 @@
 # -*- coding: utf-8 -*-
 
 """ Test SFTP client """
-import atexit
 import tempfile
 import logging
 import sys
 import os
-import random
-import string
 
 import getpass
 
@@ -49,58 +46,24 @@ def main(argv=None):  # pylint: disable=too-many-locals,too-many-statements
     log.info('Initialize Email client')
     sftp = SFTPClient(options=options, just_try=options.just_try)
     sftp.connect()
-    atexit.register(sftp.close)
 
     log.debug('Create tempory file')
-    test_content = b'Juste un test.'
+    tmp_file = tempfile.NamedTemporaryFile()  # pylint: disable=consider-using-with
-    tmp_dir = tempfile.TemporaryDirectory()  # pylint: disable=consider-using-with
+    log.debug('Temporary file path: "%s"', tmp_file.name)
-    tmp_file = os.path.join(
+    tmp_file.write(b'Juste un test.')
-        tmp_dir.name,
-        f'tmp{"".join(random.choice(string.ascii_lowercase) for i in range(8))}'
-    )
-    log.debug('Temporary file path: "%s"', tmp_file)
-    with open(tmp_file, 'wb') as file_desc:
-        file_desc.write(test_content)
-
     log.debug(
-        'Upload file %s to SFTP server (in %s)', tmp_file,
+        'Upload file %s to SFTP server (in %s)', tmp_file.name,
         options.upload_path if options.upload_path else "remote initial connection directory")
-    if not sftp.upload_file(tmp_file, options.upload_path):
+    if not sftp.upload_file(tmp_file.name, options.upload_path):
         log.error('Fail to upload test file on SFTP server')
-        sys.exit(1)
-
-    log.info('Test file uploaded on SFTP server')
-    remote_filepath = (
-        os.path.join(options.upload_path, os.path.basename(tmp_file))
-        if options.upload_path else os.path.basename(tmp_file)
-    )
-
-    with tempfile.NamedTemporaryFile() as tmp_file2:
-        log.info('Retrieve test file to %s', tmp_file2.name)
-        if not sftp.get_file(remote_filepath, tmp_file2.name):
-            log.error('Fail to retrieve test file')
-        else:
-            with open(tmp_file2.name, 'rb') as file_desc:
-                content = file_desc.read()
-                log.debug('Read content: %s', content)
-                if test_content == content:
-                    log.info('Content file retrieved match with uploaded one')
-                else:
-                    log.error('Content file retrieved doest not match with uploaded one')
-
-    try:
-        log.info('Remotly open test file %s', remote_filepath)
-        file_desc = sftp.open_file(remote_filepath)
-        content = file_desc.read()
-        log.debug('Read content: %s', content)
-        if test_content == content:
-            log.info('Content of remote file match with uploaded one')
-        else:
-            log.error('Content of remote file doest not match with uploaded one')
-    except Exception:  # pylint: disable=broad-except
-        log.exception('An exception occurred remotly opening test file %s', remote_filepath)
-
-    if sftp.remove_file(remote_filepath):
-        log.info('Test file removed on SFTP server')
     else:
-        log.error('Fail to remove test file on SFTP server')
+        log.info('Test file uploaded on SFTP server')
+        remote_filepath = (
+            os.path.join(options.upload_path, os.path.basename(tmp_file.name))
+            if options.upload_path else os.path.basename(tmp_file.name)
+        )
+        if sftp.remove_file(remote_filepath):
+            log.info('Test file removed on SFTP server')
+        else:
+            log.error('Fail to remove test file on SFTP server')
+    sftp.close()

mylib/sftp.py

@@ -101,16 +101,6 @@ class SFTPClient(ConfigurableObject):
             log.debug("Fail to retreive remote directory, use empty string instead")
             self.initial_directory = ""
 
-    def get_file(self, remote_filepath, local_filepath):
-        """ Retrieve a file from SFTP server """
-        log.debug("Retreive file '%s' to '%s'", remote_filepath, local_filepath)
-        return self.sftp_client.get(remote_filepath, local_filepath) is None
-
-    def open_file(self, remote_filepath, mode='r'):
-        """ Remotly open a file on SFTP server """
-        log.debug("Remotly open file '%s'", remote_filepath)
-        return self.sftp_client.open(remote_filepath, mode=mode)
-
     def upload_file(self, filepath, remote_directory=None):
         """ Upload a file on SFTP server """
         self.connect()
@@ -134,7 +124,8 @@ class SFTPClient(ConfigurableObject):
         if self._get_option('just_try'):
             log.debug("Just - try mode: do not really remove file '%s'", filepath)
             return True
-        return self.sftp_client.remove(filepath) is None
+        self.sftp_client.remove(filepath)
+        return True
 
     def close(self):
         """ Close SSH/SFTP connection """