Cleanups. Updated docstrings. Made pull file from remote test use sub-directories in the tree. Test cleanups.

Author: Dan

pssh/pssh_client.py

@@ -845,18 +845,9 @@ def copy_file(self, local_file, remote_file, recurse=False):
 
     def _copy_file(self, host, local_file, remote_file, recurse):
         """Make sftp client, copy file"""
-        if not host in self.host_clients or not self.host_clients[host]:
-            _user, _port, _password, _pkey = self._get_host_config_values(host)
-            self.host_clients[host] = SSHClient(
-                host, user=_user, password=_password, port=_port, pkey=_pkey,
-                forward_ssh_agent=self.forward_ssh_agent,
-                num_retries=self.num_retries,
-                timeout=self.timeout,
-                proxy_host=self.proxy_host,
-                proxy_port=self.proxy_port,
-                agent=self.agent,
-                channel_timeout=self.channel_timeout)
-        return self.host_clients[host].copy_file(local_file, remote_file, recurse=recurse)
+        self._make_ssh_client(host)
+        return self.host_clients[host].copy_file(local_file, remote_file,
+                                                 recurse=recurse)
 
     def copy_file_to_local(self, remote_file, local_file, recurse=False):
         """Copy remote file to local file in parallel
@@ -883,9 +874,19 @@ def copy_file_to_local(self, remote_file, local_file, recurse=False):
 
     def _copy_file_to_local(self, host, remote_file, local_file, recurse):
         """Make sftp client, copy file to local"""
-        if not self.host_clients[host]:
-            self.host_clients[host] = SSHClient(host, user=self.user,
-                                                password=self.password,
-                                                port=self.port, pkey=self.pkey,
-                                                forward_ssh_agent=self.forward_ssh_agent)
-        return self.host_clients[host].copy_file_to_local(remote_file, '_'.join([local_file, host]), recurse=recurse)
+        self._make_ssh_client(host)
+        return self.host_clients[host].copy_file_to_local(
+                remote_file, '_'.join([local_file, host]), recurse=recurse)
+
+    def _make_ssh_client(self, host):
+        if not host in self.host_clients or not self.host_clients[host]:
+            _user, _port, _password, _pkey = self._get_host_config_values(host)
+            self.host_clients[host] = SSHClient(
+                host, user=_user, password=_password, port=_port, pkey=_pkey,
+                forward_ssh_agent=self.forward_ssh_agent,
+                num_retries=self.num_retries,
+                timeout=self.timeout,
+                proxy_host=self.proxy_host,
+                proxy_port=self.proxy_port,
+                agent=self.agent,
+                channel_timeout=self.channel_timeout)

pssh/ssh_client.py

@@ -397,38 +397,46 @@ def copy_file_to_local(self, remote_file, local_file, recurse=False):
         :param recurse: Whether or not to recursively copy directories.
         :type recurse: bool
 
-        :raises: :mod:'ValueError' when a directory is supplied to remote_file \
+        :raises: :mod:`ValueError` when a directory is supplied to remote_file \
         and recurse is not set
+        :raises: :mod:`OSError` on OS errors creating directories or file
+        :raises: :mod:`IOError` on IO errors creating directories or file
         """
         sftp = self._make_sftp()
         try:
             sftp.listdir(remote_file)
-            remote_dir_exists = True
-        except IOError or OSError:
+        except (OSError, IOError):
             remote_dir_exists = False
+        else:
+            remote_dir_exists = True
         if remote_dir_exists and recurse:
             return self._copy_dir_to_local(remote_file, local_file)
         elif remote_dir_exists and not recurse:
             raise ValueError("Recurse must be true if remote_file is a "
                              "directory.")
         destination = self._parent_path_split(local_file)
-        if not os.path.exists(destination):
-            try:
-                os.makedirs(destination)
-            except OSError:
-                logger.error("Unable to create local directory structure.")
-                raise
+        self._make_local_dir(destination)
         try:
+            import ipdb; ipdb.set_trace()
             sftp.get(remote_file, local_file)
         except Exception, error:
             logger.error("Error occured copying file %s from remote destination %s:%s - %s",
                          local_file, self.host, remote_file, error)
+            raise error
         else:
             logger.info("Copied local file %s from remote destination %s:%s",
                         local_file, self.host, remote_file)
 
-    @staticmethod
-    def _parent_path_split(file_path):
+    def _make_local_dir(self, dirpath):
+        if not os.path.exists(dirpath):
+            try:
+                os.makedirs(dirpath)
+            except OSError:
+                logger.error("Unable to create local directory structure for "
+                             "directory %s", dirpath)
+                raise
+
+    def _parent_path_split(self, file_path):
         try:
             destination = [_dir for _dir in file_path.split(os.path.sep)
                             if _dir][:-1][0]

tests/test_ssh_client.py

@@ -185,6 +185,11 @@ def test_ssh_client_copy_remote_directory(self):
         """Tests copying a remote directory to the localhost"""
         remote_test_directory = 'remote_test_dir'
         local_test_directory = 'local_test_dir'
+        for path in [local_test_path, remote_test_path]:
+            try:
+                shutil.rmtree(path)
+            except OSError:
+                pass
         os.mkdir(remote_test_directory)
         test_files = []
         for i in range(0, 10):
@@ -263,11 +268,17 @@ def test_ssh_client_sftp_from_remote_directory(self):
         test_file_data = 'test'
         remote_test_path = 'directory_test_remote'
         local_test_path = 'directory_test_local'
+        for path in [local_test_path, remote_test_path]:
+            try:
+                shutil.rmtree(path)
+            except OSError:
+                pass
         os.mkdir(remote_test_path)
+        os.mkdir(os.path.join(remote_test_path, 'subdir'))
         local_file_paths = []
         for i in range(0, 10):
-            remote_file_path = os.path.join(remote_test_path, 'foo' + str(i))
-            local_file_path = os.path.join(local_test_path, 'foo' + str(i))
+            remote_file_path = os.path.join(remote_test_path, 'subdir', 'foo' + str(i))
+            local_file_path = os.path.join(local_test_path, 'subdir', 'foo' + str(i))
             local_file_paths.append(local_file_path)
             test_file = open(remote_file_path, 'w')
             test_file.write(test_file_data)
@@ -286,6 +297,10 @@ def test_ssh_client_remote_directory_no_recurse(self):
         test_file_data = 'test'
         remote_test_path = 'directory_test'
         local_test_path = 'directory_test_copied'
+        try:
+            shutil.rmtree(remote_test_path)
+        except OSError:
+            pass
         os.mkdir(remote_test_path)
         local_file_paths = []
         for i in range(0, 10):