SFTP client: add get_file and open_file methods

mylib/scripts/sftp_test.py

@@ -1,10 +1,13 @@
 # -*- coding: utf-8 -*-
 
 """ Test SFTP client """
+import atexit
 import tempfile
 import logging
 import sys
 import os
+import random
+import string
 
 import getpass
 
@@ -46,24 +49,58 @@ 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')
-    tmp_file = tempfile.NamedTemporaryFile()  # pylint: disable=consider-using-with
-    log.debug('Temporary file path: "%s"', tmp_file.name)
-    tmp_file.write(b'Juste un test.')
+    test_content = b'Juste un test.'
+    tmp_dir = tempfile.TemporaryDirectory()  # pylint: disable=consider-using-with
+    tmp_file = os.path.join(
+        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.name,
+        'Upload file %s to SFTP server (in %s)', tmp_file,
         options.upload_path if options.upload_path else "remote initial connection directory")
-    if not sftp.upload_file(tmp_file.name, options.upload_path):
+    if not sftp.upload_file(tmp_file, options.upload_path):
         log.error('Fail to upload test file on SFTP server')
-    else:
+        sys.exit(1)
+
     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)
+        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')
-    sftp.close()

mylib/sftp.py

@@ -101,6 +101,14 @@ 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):
+        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'):
+        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()
@@ -124,8 +132,7 @@ class SFTPClient(ConfigurableObject):
         if self._get_option('just_try'):
             log.debug("Just - try mode: do not really remove file '%s'", filepath)
             return True
-        self.sftp_client.remove(filepath)
-        return True
+        return self.sftp_client.remove(filepath) is None
 
     def close(self):
         """ Close SSH/SFTP connection """